yjxiong
diff --git a/Collapse file
‎test_models.py‎
Copy file name to clipboardExpand all lines: test_models.py
+7-1Lines changed: 7 additions & 1 deletion b/Collapse file
‎test_models.py‎
Copy file name to clipboardExpand all lines: test_models.py
+7-1Lines changed: 7 additions & 1 deletion
@@ -83,7 +83,13 @@
         batch_size=1, shuffle=False,
         num_workers=args.workers * 2, pin_memory=True)
 
-net = torch.nn.DataParallel(net, device_ids=list(range(args.workers)))
+if args.gpus is not None:
+    devices = [args.gpus[i] for i in range(args.workers)]
+else:
+    devices = list(range(args.workers))
+
+
+net = torch.nn.DataParallel(net, device_ids=devices)
 net.cuda()
 net.eval()