def test_amp_gpus_ddp_fork():
    """Ensure the use of AMP with `ddp_fork` (or associated alias strategies) does not generate CUDA initialization
    errors."""
    _ = MixedPrecision(precision="16-mixed", device="cuda")
    with multiprocessing.get_context("fork").Pool(1) as pool:
        in_bad_fork = pool.apply(torch.cuda._is_in_bad_fork)
    assert not in_bad_fork
Ensure the use of AMP with `ddp_fork` (or associated alias strategies) does not generate CUDA initialization
    errors.
tests/tests_pytorch/models/test_ddp_fork_amp.py