[fx] fixed compatiblity issue with torch 1.10 (#1331)

2025-09-22 01:48:07 +00:00 · 2022-07-18 11:41:27 +08:00
parent 069d6fdc84
commit 75abc75c15
7 changed files with 32 additions and 28 deletions
--- a/colossalai/fx/tracer/meta_patch/patched_function/arithmetic.py
+++ b/colossalai/fx/tracer/meta_patch/patched_function/arithmetic.py
@@ -3,6 +3,7 @@ from ..registry import meta_patched_function


@meta_patched_function.register(torch.matmul)
+@meta_patched_function.register('matmul')    # for built-in op @
 def torch_matmul(input, other, *, out=None):
    # copied from huggingface.utils.fx
    d1 = input.dim()