rectify gpu assignment

yjxiong · Aug 16, 2017 · bdd8742 · bdd8742
1 parent 93b3441
commit bdd8742
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/test_models.py b/test_models.py
@@ -89,8 +89,7 @@
     devices = list(range(args.workers))
 
 
-net = torch.nn.DataParallel(net, device_ids=devices)
-net.cuda()
+net = torch.nn.DataParallel(net.cuda(devices[0]), device_ids=devices)
 net.eval()
 
 data_gen = enumerate(data_loader)