Training: Only convert to cuda in cuda mode.

cmusatyalab · Jun 13, 2016 · c04386c · c04386c
1 parent f4567e8
commit c04386c
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/training/train.lua b/training/train.lua
@@ -111,9 +111,11 @@ function train()
    torch.save(paths.concat(opt.save, 'model_' .. epoch .. '.t7'), model)
    torch.save(paths.concat(opt.save, 'optimState_' .. epoch .. '.t7'), optimState)
 
-   model = model:cuda()
-   if opt.cudnn then
-      cudnn.convert(model, cudnn)
+   if opt.cuda then
+      model = model:cuda()
+      if opt.cudnn then
+         cudnn.convert(model, cudnn)
+      end
    end
    collectgarbage()
 end -- of train()