Fix "RuntimeError: Expected object of type torch.FloatTensor but foun…

…d type torch.cuda.FloatTensor for argument SeanNaren#4 'other'" for File "train.py", line 270, in <module> optimizer.step()
experimenti · Jul 23, 2018 · bcad275 · bcad275
1 parent 60ef8ba
commit bcad275
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/train.py b/train.py
@@ -139,6 +139,8 @@ def update(self, val, n=1):
  optimizer = torch.optim.SGD(parameters, lr=args.lr,
  momentum=args.momentum, nesterov=True)
  if not args.finetune: # Don't want to restart training
+ if args.cuda:
+ model.cuda()
  optimizer.load_state_dict(package['optim_dict'])
  start_epoch = int(package.get('epoch', 1)) - 1 # Index start at 0 for training
  start_iter = package.get('iteration', None)