github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-16 06:30:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx]refactor tracer (#1335)

Browse Source
This commit is contained in:
YuliangLiu0306
2022-07-19 15:50:42 +08:00
committed by GitHub
parent bf5066fba7
commit 4631fef8a0
8 changed files with 50 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_fx/test_tracer/test_timm_model/test_timm_model.py
View File

@@ -2,6 +2,7 @@ import torch
import timm.models as tm
from colossalai.fx import ColoTracer
from torch.fx import GraphModule
import pytest
def trace_and_compare(model_cls, tracer, data, meta_args=None):
@@ -53,6 +54,7 @@ def test_timm_models_without_control_flow():
trace_and_compare(model_cls, tracer, data)
@pytest.mark.skip('skip due to tracer')
def test_timm_models_with_control_flow():
torch.backends.cudnn.deterministic = True