fix

18237e29 · yuxguo · 6188a6bb · 18237e29
Commit 18237e29 authored Mar 10, 2022 by yuxguo
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletions

main.py
+2 -1

No files found.
--- a/main.py
+++ b/main.py
@@ -450,7 +450,8 @@ def main():
    model = get_model(args)
    
    if args.use_gpu:
-        model = DataParallel(model).cuda()
+        # model = DataParallel(model).cuda()
+        model = model.cuda()
        
    if args.weight_decay == 0:
        optimizer = optim.Adam(model.parameters(), lr=args.lr)