From 5875fadc67f698f9ab4d707ccda5ede0c203b9ca Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 07:26:18 -0400
Subject: [PATCH] added cpu model test

---
 .../new_project_templates/lightning_module_template.py        | 2 +-
 pytorch_lightning/models/trainer.py                           | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/pytorch_lightning/examples/new_project_templates/lightning_module_template.py b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
index 7f5459e1..df7dbe32 100644
--- a/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
+++ b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
@@ -108,7 +108,7 @@ class LightningTemplateModel(LightningModule):
 
         output = OrderedDict({
             'val_loss': loss_val,
-            'val_acc': torch.tensor(val_acc).cuda(loss_val.device.index),
+            'val_acc': torch.tensor(val_acc).type(loss_val.dtype),
         })
 
         # can also return just a scalar instead of a dict (return loss_val)
diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index fea7213e..b7f8fbe7 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -436,6 +436,10 @@ class Trainer(TrainerIO):
 
             self.__run_pretrain_routine(model)
 
+        # return 1 when finished
+        # used for testing or when we need to know that training succeeded
+        return 1
+
     def dp_train(self, model):
 
         # CHOOSE OPTIMIZER