Fix issues in lightning optimizer refactor merge

Lightning-AI · Jan 12, 2021 · 32cba95 · 32cba95
1 parent a9b30cd
commit 32cba95
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/pytorch_lightning/core/optimizer.py b/pytorch_lightning/core/optimizer.py
@@ -17,7 +17,7 @@
 
 from torch.optim.optimizer import Optimizer
 
-from pytorch_lightning.utilities import _TPU_AVAILABLE, DeviceType
+from pytorch_lightning.utilities import _TPU_AVAILABLE, DeviceType, AMPType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 if _TPU_AVAILABLE:

diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
@@ -31,6 +31,7 @@
 from pytorch_lightning.utilities.model_helpers import is_overridden
 from pytorch_lightning.utilities.parsing import AttributeDict
 from pytorch_lightning.utilities.warnings import WarningCache
+from pytorch_lightning.core.optimizer import LightningOptimizer
 
 
 class TrainLoop: