From 19227238e670f8631dffb0248c561f3f49ab0d4c Mon Sep 17 00:00:00 2001 From: Ivan Shivalov <142020129+intexcor@users.noreply.github.com> Date: Thu, 12 Dec 2024 13:37:07 +0300 Subject: [PATCH] add LTSM Layer support add LTSM Layer support --- torchsummary/torchsummary.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 1ed065f..0419b9b 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, torch.Tensor) ] else: summary[m_key]["output_shape"] = list(output.size())