Merge pull request #515 from AlexanderHustinx/pytorch-1.0

solve error when resuming training
jwyang · Apr 18, 2019 · aec4244 · aec4244
2 parents 297e8e0 + e9ce564
commit aec4244
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/trainval_net.py b/trainval_net.py
@@ -262,6 +262,9 @@ def __len__(self):
       else:
         params += [{'params':[value],'lr':lr, 'weight_decay': cfg.TRAIN.WEIGHT_DECAY}]
 
+  if args.cuda:
+    fasterRCNN.cuda()
+
   if args.optimizer == "adam":
     lr = lr * 0.1
     optimizer = torch.optim.Adam(params)
@@ -286,9 +289,6 @@ def __len__(self):
   if args.mGPUs:
     fasterRCNN = nn.DataParallel(fasterRCNN)
 
-  if args.cuda:
-    fasterRCNN.cuda()
-
   iters_per_epoch = int(train_size / args.batch_size)
 
   if args.use_tfboard: