added eval and train for redundancy (#464)

2019-11-05 09:14:33 -05:00 · 2019-11-05 09:14:33 -05:00 · 8fbaccddae
parent 32dd803b1e
commit 8fbaccddae
1 changed files with 4 additions and 0 deletions
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@ -247,6 +247,10 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
        for param in self.parameters():
            param.requires_grad = False

+        self.eval()
+
    def unfreeze(self):
        for param in self.parameters():
            param.requires_grad = True
+
+        self.train()