Skip to content

GitLab

Commit fbd2a9e0 authored by YuliangLiu0306's avatar YuliangLiu0306 Committed by アマデウス
Browse files

[hotfix] meta_tensor_compatibility_with_torch2

parent 15a74da7
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 3 deletions
colossalai/fx/_compatibility.py
View file @ fbd2a9e0
...@@ -14,9 +14,7 @@ elif TORCH_MAJOR == 1 and TORCH_MINOR == 13: ...@@ -14,9 +14,7 @@ elif TORCH_MAJOR == 1 and TORCH_MINOR == 13:
from . import _meta_regist_13 from . import _meta_regist_13
META_COMPATIBILITY = True META_COMPATIBILITY = True
elif TORCH_MAJOR == 2: elif TORCH_MAJOR == 2:
from . import _meta_regist_13
META_COMPATIBILITY = True META_COMPATIBILITY = True
raise UserWarning("Colossalai is not tested with torch2.0 yet!!!")
def compatibility(is_backward_compatible: bool = False) -> Callable: def compatibility(is_backward_compatible: bool = False) -> Callable:
......
colossalai/fx/profiler/opcount.py
View file @ fbd2a9e0
...@@ -223,7 +223,8 @@ def zero_flop_jit(*args): ...@@ -223,7 +223,8 @@ def zero_flop_jit(*args):
return 0 return 0
if version.parse(torch.__version__) >= version.parse('1.12.0'): if version.parse(torch.__version__) >= version.parse('1.12.0') and version.parse(
torch.__version__) < version.parse('2.0.0'):
flop_mapping = { flop_mapping = {
# gemm, gemv and dot # gemm, gemv and dot
aten.mm.default: matmul_flop_jit, aten.mm.default: matmul_flop_jit,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment