Update main.py

Using cuda.device_count instead of hardcoding the number of GPUs.

Author: fducau

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()