From 6169d22813fb50d594fd543a0bd18cfda5ca7008 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 24 Jul 2019 13:59:51 -0400 Subject: [PATCH] refactored model tests --- .../new_project_templates/lightning_module_template.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 59d7210a..c8622230 100644 --- a/pytorch_lightning/examples/new_project_templates/lightning_module_template.py +++ b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py @@ -85,7 +85,7 @@ class LightningTemplateModel(LightningModule): loss_val = self.loss(y, y_hat) output = OrderedDict({ - 'loss': loss_val + 'loss': loss_val.unsqueeze(1) }) # can also return just a scalar instead of a dict (return loss_val) @@ -112,8 +112,8 @@ class LightningTemplateModel(LightningModule): val_acc = val_acc.cuda(loss_val.device.index) output = OrderedDict({ - 'val_loss': loss_val, - 'val_acc': val_acc, + 'val_loss': loss_val.unsqueeze(1), + 'val_acc': val_acc.unsqueeze(1), }) # can also return just a scalar instead of a dict (return loss_val)