Fix amp autocast

pytorch_lightning/overrides/base.py

@@ -19,12 +19,13 @@
 
 from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.trainer.states import RunningStage
+from pytorch_lightning.utilities.device_dtype_mixin import DeviceDtypeModuleMixin
 from pytorch_lightning.utilities.warnings import WarningCache
 
 warning_cache = WarningCache()
 
 
-class _LightningModuleWrapperBase(torch.nn.Module):
+class _LightningModuleWrapperBase(DeviceDtypeModuleMixin, torch.nn.Module):
 
     def __init__(self, pl_module: LightningModule):
         """

pytorch_lightning/plugins/precision/native_amp.py

@@ -91,4 +91,5 @@ def post_optimizer_step(self, optimizer: Optimizer, optimizer_idx: int) -> None:
     @contextmanager
     def train_step_context(self) -> Generator[autocast, None, None]:
         """Enable autocast context"""
-        yield torch.cuda.amp.autocast()
+        with torch.cuda.amp.autocast():
+            yield

tests/models/test_amp.py

@@ -63,7 +63,7 @@ def test_amp_single_gpu_ddp_spawn(tmpdir):
     model = BoringModel()
     # tutils.run_model_test(trainer_options, model)
     trainer.fit(model)
-
+    assert torch.is_autocast_enabled()
     assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
 
 
@@ -103,7 +103,7 @@ def test_amp_multi_gpu_ddp_spawn(tmpdir):
     model = BoringModel()
     # tutils.run_model_test(trainer_options, model)
     trainer.fit(model)
-
+    assert torch.is_autocast_enabled()
     assert trainer.state == TrainerState.FINISHED, f"Training failed with {trainer.state}"
 
 
@@ -152,6 +152,7 @@ def test_amp_gpu_ddp_slurm_managed(tmpdir):
     assert trainer.training_type_plugin.cluster_environment.resolve_root_node_address('abc[23-24]') == 'abc23'
     generated = trainer.training_type_plugin.cluster_environment.resolve_root_node_address('abc[23-24, 45-40, 40]')
     assert generated == 'abc23'
+    assert torch.is_autocast_enabled()
 
 
 @pytest.mark.skipif(torch.cuda.is_available(), reason="test is restricted only on CPU")