broadcast the two input shapes for transposed matmul

pytorch · Oct 19, 2022 · bfa971a · bfa971a
1 parent a9a4bb2
commit bfa971a
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/py/torch_tensorrt/fx/converters/acc_ops_converters.py b/py/torch_tensorrt/fx/converters/acc_ops_converters.py
@@ -37,6 +37,14 @@ def trt_transposed_matmul_converter(network, target, args, kwargs, name):
         lhs = get_trt_tensor(network, lhs, f"{name}_lhs")
     if isinstance(rhs, torch.nn.Parameter):
         rhs = get_trt_tensor(network, rhs, f"{name}_rhs")
+
+    lhs, rhs = broadcast(
+        network,
+        lhs,
+        rhs,
+        f"{lhs.name}_broadcast",
+        f"{rhs.name}_broadcast",
+    )
     layer = network.add_matrix_multiply(
         lhs,
         trt.MatrixOperation.TRANSPOSE if lhs_transposed else trt.MatrixOperation.NONE,