Fix handling input_size with multi-input

Author: ahmedhshahin

torchsummary/torchsummary.py

@@ -98,7 +98,9 @@ def hook(module, input, output):
         summary_str += line_new + "\n"
 
     # assume 4 bytes/number (float on cuda).
-    total_input_size = abs(np.prod(sum(input_size, ()))
+    # to handle the case of multi-input: prod(input1) + prod(input2) + ...
+    n_input_size = np.array([np.prod(i) for i in input_size]).sum() if isinstance(input_size, list) else np.prod(input_size)
+    total_input_size = abs(n_input_size
                            * batch_size * 4. / (1024 ** 2.))
     total_output_size = abs(2. * total_output * 4. /
                             (1024 ** 2.))  # x2 for gradients