This commit is contained in:
Glenn Jocher 2018-12-28 20:11:10 +01:00
parent eec0dc7b6c
commit 16bc3b72c3
1 changed files with 1 additions and 0 deletions

View File

@ -13,6 +13,7 @@ def init_seeds(seed=0):
if CUDA_AVAILABLE: if CUDA_AVAILABLE:
torch.cuda.manual_seed(seed) torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed) torch.cuda.manual_seed_all(seed)
# torch.cuda.set_device(0) # OPTIONAL: Set your GPU if multiple available
def select_device(force_cpu=False): def select_device(force_cpu=False):