added sample input for summary

2019-07-24 16:28:55 -04:00 · 2019-07-24 16:28:55 -04:00 · 8db8cd2539
parent 77a7f3e33e
commit 8db8cd2539
1 changed files with 1 additions and 1 deletions
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@ -41,7 +41,7 @@ class ModelSummary(object):
        if self.model.trainer.use_amp:
            input_ = input_.half()

-        with torch.no_grad:
+        with torch.no_grad():

            for i in range(1, len(mods)):
                m = mods[i]