Merge pull request jwyang#327 from JakeCowton/gpu-resume

Move model to GPU before loading checkpoint data
sanjin0000 · Oct 8, 2018 · bbef869 · bbef869
2 parents 3093199 + b4d4c52
commit bbef869
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/trainval_net.py b/trainval_net.py
@@ -269,6 +269,9 @@ def __len__(self):
   elif args.optimizer == "sgd":
     optimizer = torch.optim.SGD(params, momentum=cfg.TRAIN.MOMENTUM)
 
+  if args.cuda:
+    fasterRCNN.cuda()
+
   if args.resume:
     load_name = os.path.join(output_dir,
       'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint))
@@ -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: