From fd84d8e581dbccd4255713abea48e14570c9265d Mon Sep 17 00:00:00 2001 From: wyc1997 Date: Sat, 22 Feb 2020 14:36:50 -0800 Subject: [PATCH] add checks for atomic values in output --- torchsummary/torchsummary.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 1ed065f..e18aa85 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -31,7 +31,7 @@ def hook(module, input, output): summary[m_key]["input_shape"][0] = batch_size if isinstance(output, (list, tuple)): summary[m_key]["output_shape"] = [ - [-1] + list(o.size())[1:] for o in output + [-1] + list(o.size())[1:] for o in output if isinstance(o, (list, tuple)) ] else: summary[m_key]["output_shape"] = list(output.size())