respect multigpu envs, e.g. cuda:2 designation should work

eliplutchok · Aug 19, 2022 · 055e7ee · 055e7ee
1 parent 013af92
commit 055e7ee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/makemore.py b/makemore.py
@@ -411,7 +411,7 @@ def create_datasets(input_file):
         loss.backward()
         torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
         optimizer.step()
-        if args.device == 'cuda':
+        if args.device.startswith('cuda'):
             torch.cuda.synchronize()
         t1 = time.time()