added sample input for summary

2019-07-24 16:24:58 -04:00 · 2019-07-24 16:24:58 -04:00 · b8cc62ee52
parent 3a86e0fc6c
commit b8cc62ee52
1 changed files with 5 additions and 0 deletions
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@ -34,6 +34,11 @@ class ModelSummary(object):
        in_sizes = []
        out_sizes = []
        input_ = self.model.example_input_array
+
+        if self.model.on_gpu:
+            input_ = input_.cuda(0)
+
+
        for i in range(1, len(mods)):
            m = mods[i]
            if type(input_) is list or type(input_) is tuple: