utils.py 293 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
import torch

dtypes = [torch.float, torch.double]

devices = [torch.device('cpu')]
rusty1s's avatar
0.4.1  
rusty1s committed
6
if torch.cuda.is_available():
rusty1s's avatar
rusty1s committed
7
    devices += [torch.device(f'cuda:{torch.cuda.current_device()}')]
rusty1s's avatar
rusty1s committed
8
9
10
11


def tensor(x, dtype, device):
    return None if x is None else torch.tensor(x, dtype=dtype, device=device)