github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

[test] refactor tests with spawn (#3452)

Browse Source 
* [test] added spawn decorator

* polish code

* polish code

* polish code

* polish code

* polish code

* polish code
This commit is contained in:
Frank Lee
2023-04-06 14:51:35 +08:00
committed by GitHub
parent 62f4e2eb07
commit 80eba05b0a
240 changed files with 1723 additions and 2342 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tests/test_fx/test_tracer/test_patched_op.py
View File

@@ -1,7 +1,10 @@
import torch
from colossalai.fx.tracer.meta_patch import patched_function
from functools import partial
import torch
from colossalai.fx.tracer.meta_patch import patched_function
from colossalai.testing import clear_cache_before_run
def _run(data, patch_fn):
try:
@@ -22,6 +25,7 @@ def _assert_output_shape(data, patch_fn, expect_exception, output_shape):
assert output.shape == output_shape
@clear_cache_before_run()
def test_repeat_interleave():
patch_fn = patched_function.torch_repeat_interleave
@@ -63,6 +67,7 @@ def test_repeat_interleave():
output_shape=materialized_output.shape)
@clear_cache_before_run()
def test_torch_max():
data = torch.rand(4, 3)
out = torch.max(data)