Update Readme so that .test will work. (#659)

When one follows the Readme, the example will fail once we call `trainer.test()` because the methods are not overridden.

Fixes https://github.com/williamFalcon/pytorch-lightning/issues/428
This commit is contained in:
Frédéric Branchaud-Charron
2020-01-13 22:27:53 -05:00
committed by William Falcon
parent ec7fc97857
commit 083dd6a3ef
+12
View File
@@ -128,6 +128,18 @@ To use lightning do 2 things:
avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
tensorboard_logs = {'val_loss': avg_loss}
return {'avg_val_loss': avg_loss, 'log': tensorboard_logs}
def test_step(self, batch, batch_idx):
# OPTIONAL
x, y = batch
y_hat = self.forward(x)
return {'test_loss': F.cross_entropy(y_hat, y)}
def test_end(self, outputs):
# OPTIONAL
avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
tensorboard_logs = {'test_loss': avg_loss}
return {'avg_test_loss': avg_loss, 'log': tensorboard_logs}
def configure_optimizers(self):
# REQUIRED