Fix summary bug when calaculating output shape (#31549)

* fix summary bug
4 years ago · 75433126df
parent c3634c6b0a
commit 75433126df
1 changed files with 4 additions and 2 deletions
--- a/python/paddle/hapi/model_summary.py
+++ b/python/paddle/hapi/model_summary.py
@ -341,10 +341,12 @@ def summary_string(model, input_size, dtypes=None):
        total_params += summary[layer]["nb_params"]

        try:
-            total_output += np.prod(summary[layer]["output_shape"])
+            total_output += np.sum(
+                np.prod(
+                    summary[layer]["output_shape"], axis=-1))
        except:
            for output_shape in summary[layer]["output_shape"]:
-                total_output += np.prod(output_shape)
+                total_output += np.sum(np.prod(output_shape, axis=-1))

        if "trainable" in summary[layer]:
            if summary[layer]["trainable"] == True: