Fixes

Author: carmocca

tests/tests_pytorch/models/test_ddp_fork_amp.py

@@ -15,7 +15,7 @@
 
 import torch
 
-from pytorch_lightning.plugins import NativeMixedPrecisionPlugin
+from pytorch_lightning.plugins import MixedPrecisionPlugin
 from tests_pytorch.helpers.runif import RunIf
 
 
@@ -24,7 +24,7 @@
 def test_amp_gpus_ddp_fork():
     """Ensure the use of native AMP with `ddp_fork` (or associated alias strategies) does not generate CUDA
     initialization errors."""
-    _ = NativeMixedPrecisionPlugin(precision=16, device="cuda")
+    _ = MixedPrecisionPlugin(precision=16, 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