src.acoustools.Utilities.Setup

 1import torch, sys
 2from typing import Literal 
 3
 4DTYPE = torch.complex64
 5DTYPE = torch.complex32 if '-c32' in sys.argv else DTYPE
 6DTYPE = torch.complex64 if '-c64' in sys.argv else DTYPE
 7DTYPE = torch.complex128 if '-c128' in sys.argv else DTYPE
 8'''
 9Data type to use for matricies - use `.to(DTYPE)` to convert
10'''
11
12device:Literal['cuda','cpu'] = 'cuda' if torch.cuda.is_available() else 'cpu' 
13# device = 'mps' if torch.backends.mps.is_available() else device
14'''Constant storing device to use, `cuda` if cuda is available else cpu. \n
15Use -cpu when running python to force cpu use'''
16device = device if '-cpu' not in sys.argv else 'cpu'
17device = device if '-mps' not in sys.argv else 'mps'