From d004fc57258e3aa8d67a305abb7c9ca53a25e271 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 24 Jul 2019 14:00:29 -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 c8622230..10e6c028 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.unsqueeze(1) + 'loss': loss_val.unsqueeze(0) }) # 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.unsqueeze(1), - 'val_acc': val_acc.unsqueeze(1), + 'val_loss': loss_val.unsqueeze(0), + 'val_acc': val_acc.unsqueeze(0), }) # can also return just a scalar instead of a dict (return loss_val)