cuda set device

shubhamagarwal92 · shubhamagarwal92 · commit 6c4e2b864186 · 2020-05-29T00:06:44.000+01:00
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))

-Original file line number
+Diff line change
 # virtualenv
 venv/
 ENV/
++
 +.idea/
Original file line number	Diff line number	Diff line change
`@@ -113,6 +113,7 @@`
`113`	`113`	`if args.gpu_ids[0] >= 0`
`114`	`114`	`else torch.device("cpu")`
`115`	`115`	`)`
	`116`	`+torch.cuda.set_device(device)`
`116`	`117`
`117`	`118`	`# Print config and args.`
`118`	`119`	`print(yaml.dump(config, default_flow_style=False))`