github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-11 05:49:55 +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

4
tests/test_checkpoint_io/test_general_checkpoint_io.py
View File

@@ -6,6 +6,7 @@ from torch.optim import Adam
from torchvision.models import resnet18
from colossalai.checkpoint_io import GeneralCheckpointIO
from colossalai.testing import clear_cache_before_run, parameterize
# ========
# Note:
@@ -15,7 +16,8 @@ from colossalai.checkpoint_io import GeneralCheckpointIO
# ========
@pytest.mark.parametrize('use_safetensors', [True, False])
@clear_cache_before_run()
@parameterize('use_safetensors', [True, False])
def test_unsharded_checkpoint(use_safetensors: bool):
# create a model and optimizer
model = resnet18()