Merge pull request #1120 from pytorch/fb-sync-khabinov

[fx_acc] Add acc_tracer support for torch.mm
pytorch · Jun 15, 2022 · c558de8 · c558de8
2 parents 058a511 + d897a0a
commit c558de8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/docs/v1.1.0/._index.html b/docs/v1.1.0/._index.html
diff --git a/py/torch_tensorrt/fx/tracer/acc_tracer/acc_ops.py b/py/torch_tensorrt/fx/tracer/acc_tracer/acc_ops.py
@@ -739,6 +739,13 @@ def square_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
         ("mat2", "other"),
     ],
 )
+@register_acc_op_mapping(
+    op_and_target=("call_function", torch.mm),
+    arg_replacement_tuples=[
+        ("input", "input"),
+        ("mat2", "other"),
+    ],
+)
 @register_acc_op_mapping(op_and_target=("call_function", torch.matmul))
 @register_acc_op
 def matmul(*, input, other):