24 lines
447 B
Python
24 lines
447 B
Python
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def check_cuda():
|
||
|
|
return torch.cuda.is_available()
|
||
|
|
|
||
|
|
|
||
|
|
CUDA_AVAILABLE = check_cuda()
|
||
|
|
|
||
|
|
|
||
|
|
def init_seeds(seed=0):
|
||
|
|
torch.manual_seed(seed)
|
||
|
|
if CUDA_AVAILABLE:
|
||
|
|
torch.cuda.manual_seed(seed)
|
||
|
|
torch.cuda.manual_seed_all(seed)
|
||
|
|
|
||
|
|
|
||
|
|
def select_device(force_cpu=False):
|
||
|
|
if force_cpu:
|
||
|
|
device = torch.device('cpu')
|
||
|
|
else:
|
||
|
|
device = torch.device('cuda:0' if CUDA_AVAILABLE else 'cpu')
|
||
|
|
return device
|