[fix]

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

Showing with 3 additions and 2 deletions

main.py main.py +3 -2

No files found.
--- a/main.py
+++ b/main.py
@@ -214,8 +214,9 @@ def main():
    cudnn.benchmark = True
    net.cuda()
    device_ids = [0, 1, 2, 3, 4]
-    net = nn.DataParallel(net, device_ids=device_ids)
-    optimizer = nn.DataParallel(optimizer, 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:
        #net.load_state_dict(torch.load(checkpath))