added sample input for summary

2026-09-10 12:21:57 +08:00 · 2019-07-24 16:24:58 -04:00
parent 3a86e0fc6c
commit b8cc62ee52
1 changed files with 5 additions and 0 deletions
@@ -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: