added fallback local init

2019-07-13 10:16:50 -04:00 · 2019-07-13 10:16:50 -04:00 · 52a98d76d8
parent 8b0cda84e7
commit 52a98d76d8
1 changed files with 1 additions and 1 deletions
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@ -597,7 +597,7 @@ class Trainer(TrainerIO):
        try:
            loss = output['loss']
        except Exception as e:
-            if type(loss) is torch.Tensor:
+            if type(output) is torch.Tensor:
                loss = output

        self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)