[fix]

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

Showing with 3 additions and 1 deletion

main.py main.py +3 -1

No files found.
--- a/main.py
+++ b/main.py
@@ -213,7 +213,9 @@ def main():

    cudnn.benchmark = True
    net.cuda()
-    net = nn.DataParallel(net, device_ids=[0, 1, 2, 3, 4])
+    device_ids = [0, 1, 2, 3, 4]
+    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))