refactored model tests

This commit is contained in:
William Falcon
2019-07-24 14:00:29 -04:00
parent 6169d22813
commit d004fc5725
@@ -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)