From b4d4c5246898e62ac957ce949b025a8bf64cd599 Mon Sep 17 00:00:00 2001
From: Jake Cowton <j.cowton2@ncl.ac.uk>
Date: Wed, 3 Oct 2018 13:46:02 +0100
Subject: [PATCH] Move model to GPU before loading checkpoint data

---
 trainval_net.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/trainval_net.py b/trainval_net.py
index ffc929f41..2086c35e9 100644
--- 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: