Merge pull request pytorch#3 from fducau/master

kuangliu · web-flow · commit 1a4bac9aa5e7 · 2017-04-28T10:06:14.000+08:00
Solve not working with 1 GPU
diff --git a/main.py b/main.py
@@ -63,7 +63,7 @@
 
 if use_cuda:
     net.cuda()
-    net = torch.nn.DataParallel(net, device_ids=[0,1,2,3])
+    net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
     cudnn.benchmark = True
 
 criterion = nn.CrossEntropyLoss()
