[PIR] Fix fc api bug in pir mode (#68806)

zhangbo9674 · web-flow · commit c9181135a655 · 2024-10-21T10:21:46.000+08:00
* fix bug

* fix
diff --git a/python/paddle/static/nn/common.py b/python/paddle/static/nn/common.py
@@ -220,10 +220,9 @@ def fc_base(
             )
             if in_pir_mode():
                 if len(input_var.shape) > 2:
-                    new_shape = (
-                        input_var.shape[0],
-                        np.prod(input_var.shape[1:]),
-                    )
+                    new_shape = input_var.shape[:num_flatten_dims] + [
+                        np.prod(input_var.shape[num_flatten_dims:])
+                    ]
                     input_var = paddle.reshape(input_var, new_shape)
                 tmp = paddle.matmul(input_var, w)
             else:
