diff --git a/tests/test_booster/test_mixed_precision/test_fp16_torch.py b/tests/test_booster/test_mixed_precision/test_fp16_torch.py index 808b11d87..3fd6b7df1 100644 --- a/tests/test_booster/test_mixed_precision/test_fp16_torch.py +++ b/tests/test_booster/test_mixed_precision/test_fp16_torch.py @@ -3,11 +3,10 @@ from torch.optim import Adam import colossalai from colossalai.booster.mixed_precision import FP16TorchMixedPrecision -from colossalai.testing import clear_cache_before_run, rerun_if_address_is_in_use, spawn +from colossalai.testing import rerun_if_address_is_in_use, spawn from tests.kit.model_zoo import model_zoo -@clear_cache_before_run() def run_torch_amp(rank, world_size, port): # init dist env colossalai.launch(rank=rank, world_size=world_size, port=port, host="localhost") @@ -35,7 +34,6 @@ def run_torch_amp(rank, world_size, port): optimizer.step() del model, optimizer, criterion, data, output, mixed_precision - @rerun_if_address_is_in_use() def test_torch_ddp_plugin(): spawn(run_torch_amp, 1)