diff --git a/torch_utils/misc.py b/torch_utils/misc.py
index 7829f4d9f..8732d7437 100755
--- a/torch_utils/misc.py
+++ b/torch_utils/misc.py
@@ -235,7 +235,7 @@ def post_hook(mod, _inputs, outputs):
         name = '<top-level>' if e.mod is module else submodule_names[e.mod]
         param_size = sum(t.numel() for t in e.unique_params)
         buffer_size = sum(t.numel() for t in e.unique_buffers)
-        output_shapes = [str(list(e.outputs[0].shape)) for t in e.outputs]
+        output_shapes = [str(list(t.shape)) for t in e.outputs]
         output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs]
         rows += [[
             name + (':0' if len(e.outputs) >= 2 else ''),