[fix]

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

Showing with 1 addition and 1 deletion

main.py main.py +1 -1

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

    cudnn.benchmark = True
    net.cuda()
-    net = nn.DataParallel(net)
+    net = nn.DataParallel(net, device_ids=[0, 1, 2, 3, 4])

    if args.use_checkpoint:
        #net.load_state_dict(torch.load(checkpath))