This commit is contained in:
flybird11111 2025-04-01 15:32:11 +08:00
parent 822556a8ca
commit 4b8b67ae23

View File

@ -1,5 +1,6 @@
import torch
from torch.optim import Adam
import pytest
import colossalai
from colossalai.booster.mixed_precision import FP16TorchMixedPrecision
@ -35,6 +36,7 @@ def run_torch_amp(rank, world_size, port):
del model, optimizer, criterion, data, output, mixed_precision
@pytest.mark.skip("test ci.")
@rerun_if_address_is_in_use()
def test_torch_ddp_plugin():
spawn(run_torch_amp, 1)