From 083dd6a3ef39f5585a91b15cbd178aa7d9ce1e90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Branchaud-Charron?= Date: Mon, 13 Jan 2020 22:27:53 -0500 Subject: [PATCH] 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 --- README.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/README.md b/README.md index 84e44586..e91522dc 100644 --- a/README.md +++ b/README.md @@ -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