Update flops.py

yghstill · web-flow · commit 874bba39a2b0 · 2023-04-17T17:33:35.000+08:00
diff --git a/paddleslim/analysis/flops.py b/paddleslim/analysis/flops.py
@@ -86,7 +86,7 @@ def _graph_flops(graph, only_conv=True, detail=False):
             k_size = op.attr("ksize")
             flops += h_out * w_out * c_out * (k_size[0]**2)
 
-        elif op.type() == 'mul':
+        elif op.type() in ['mul', 'matmul', 'matmul_v2']:
             x_shape = list(op.inputs("X")[0].shape())
             y_shape = op.inputs("Y")[0].shape()
             if x_shape[0] == -1:
