cuda set device

batra-mlp-lab · May 28, 2020 · 6c4e2b8 · 6c4e2b8
1 parent 556def3
commit 6c4e2b8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -37,3 +37,5 @@ logs/
 # virtualenv
 venv/
 ENV/
+
+.idea/
diff --git a/train.py b/train.py
@@ -113,6 +113,7 @@
     if args.gpu_ids[0] >= 0
     else torch.device("cpu")
 )
+torch.cuda.set_device(device)
 
 # Print config and args.
 print(yaml.dump(config, default_flow_style=False))