github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-15 22:19:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] meta_tensor_compatibility_with_torch2

Browse Source
This commit is contained in:
YuliangLiu0306
2023-03-30 11:22:20 +08:00
committed by アマデウス
parent 15a74da79c
commit fbd2a9e05b
2 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/fx/profiler/opcount.py
View File

@@ -223,7 +223,8 @@ def zero_flop_jit(*args):
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 = {
# gemm, gemv and dot
aten.mm.default: matmul_flop_jit,