utils.py 366 Bytes
Newer Older
rusty1s's avatar
rusty1s committed
1
import torch
rusty1s's avatar
rusty1s committed
2

rusty1s's avatar
rusty1s committed
3
dtypes = [torch.float, torch.double, torch.int, torch.long]
rusty1s's avatar
rusty1s committed
4
grad_dtypes = [torch.float, torch.double]
rusty1s's avatar
rusty1s committed
5

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

rusty1s's avatar
rusty1s committed
10
11
12

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