[fix]

b2c0db9b · lishen · 07fb05e0 · b2c0db9b
Commit b2c0db9b authored May 18, 2023 by lishen
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

main.py main.py +2 -2

No files found.
--- a/main.py
+++ b/main.py
@@ -214,8 +214,8 @@ def main():
    cudnn.benchmark = True
    net.cuda()
    device_ids = [0, 1, 2, 3, 4]
-    net = nn.DataParallel(net).to(device_ids)
-    # net = nn.DataParallel(net, device_ids=device_ids)
+    # net = nn.DataParallel(net).to(device_ids)
+    net = nn.DataParallel(net, device_ids=device_ids)
    # optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

    if args.use_checkpoint: