fix

40d3539c · yuxguo · 67498b6b · 40d3539c
Commit 40d3539c authored Mar 14, 2022 by yuxguo
Hide whitespace changes
Inline Side-by-side

Showing with 1 additions and 1 deletions

main.py
+1 -1

No files found.
--- a/main.py
+++ b/main.py
@@ -454,7 +454,7 @@ def main():
    model = get_model(args)
    
    if args.use_gpu:
-        model = DataParallel(model).cuda()
+        model = DataParallel(model, device_ids=[0, 1]).cuda()
        # model = model.cuda()
        
    if args.weight_decay == 0: