Use torch.autocast (#10053)

Author: carmocca

CHANGELOG.md

@@ -206,6 +206,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
   * Added bfloat16 support for Lightning Trainer ([#9049](https://github.com/PyTorchLightning/pytorch-lightning/pull/9049))
   * Renamed `TPUHalfPrecisionPlugin` to `TPUBf16PrecisionPlugin` ([#10026](https://github.com/PyTorchLightning/pytorch-lightning/pull/10026))
   * Default to `precision=bf16` on CPU when `precision=16` is passed ([#10033](https://github.com/PyTorchLightning/pytorch-lightning/pull/10033))
+  * Add support for `torch.autocast` ([#10053](https://github.com/PyTorchLightning/pytorch-lightning/pull/10053))
 
 
 - Added `kfold` example for loop customization ([#9965](https://github.com/PyTorchLightning/pytorch-lightning/pull/9965))

pytorch_lightning/plugins/precision/native_amp.py

@@ -24,6 +24,11 @@
 from pytorch_lightning.utilities import _TORCH_GREATER_EQUAL_DEV_1_10, AMPType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
+if _TORCH_GREATER_EQUAL_DEV_1_10:
+    from torch import autocast
+else:
+    from torch.cuda.amp import autocast
+
 
 class NativeMixedPrecisionPlugin(MixedPrecisionPlugin):
     """Plugin for native mixed precision training with :mod:`torch.cuda.amp`.
@@ -90,12 +95,10 @@ def pre_optimizer_step(
             self.scaler.update()
         return False
 
-    def autocast_context_manager(self) -> torch.cuda.amp.autocast:
-        if self.use_cpu:
-            return torch.cpu.amp.autocast(dtype=self._dtype)  # Only reached in pytorch==1.10 where this is ok. skipcq
-        if self.is_bfloat16:
-            return torch.cuda.amp.autocast(dtype=self._dtype)  # Only reached in pytorch==1.10 where this is ok. skipcq
-        return torch.cuda.amp.autocast()
+    def autocast_context_manager(self) -> autocast:
+        if _TORCH_GREATER_EQUAL_DEV_1_10:
+            return autocast("cpu" if self.use_cpu else "cuda", dtype=self._dtype)
+        return autocast()
 
     @contextmanager
     def forward_context(self) -> Generator[None, None, None]:

tests/plugins/test_amp_plugins.py

@@ -181,12 +181,11 @@ def test_amp_apex_ddp_spawn_fit(amp_level, tmpdir):
 @pytest.mark.skipif(not _TORCH_GREATER_EQUAL_DEV_1_10, reason="Torch CPU AMP is not available.")
 def test_cpu_amp_precision_context_manager(tmpdir):
     """Test to ensure that the context manager correctly is set to CPU + bfloat16, and a scaler isn't set."""
-
     plugin = NativeMixedPrecisionPlugin(precision="bf16", use_cpu=True)
     assert plugin.use_cpu
     assert not hasattr(plugin, "scaler")
     context_manager = plugin.autocast_context_manager()
-    assert isinstance(context_manager, torch.cpu.amp.autocast)
+    assert isinstance(context_manager, torch.autocast)
     assert context_manager.fast_dtype == torch.bfloat16