[Training] Fix lightning _PATH import

ShaohanTian · Mar 29, 2023 · 009a3e7 · 009a3e7
1 parent 993d124
commit 009a3e7
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/training/src/utils/ddp_zero1.py b/training/src/utils/ddp_zero1.py
@@ -9,8 +9,13 @@
 
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.core.optimizer import LightningOptimizer
-from pytorch_lightning.utilities.types import _PATH
-# from lightning_lite.utilities.types import _PATH
+try:  # pytorch_lightning <= 1.7
+    from pytorch_lightning.utilities.types import _PATH
+except ImportError:  # pytorch_lightning >= 1.8
+    try:
+        from lightning_lite.utilities.types import _PATH
+    except ImportError:  # pytorch_lightning >= 1.9
+        from lightning_fabric.utilities.types import _PATH
 
 
 # Copied from Pytorch's ZeroRedundancyOptimizer's state_dict method, but we only get

diff --git a/training/src/utils/ddp_zero2.py b/training/src/utils/ddp_zero2.py
@@ -13,9 +13,14 @@
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.plugins.precision import PrecisionPlugin, NativeMixedPrecisionPlugin
 from pytorch_lightning.core.optimizer import LightningOptimizer
-from pytorch_lightning.utilities.types import _PATH
-# from lightning_lite.utilities.types import _PATH
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+try:  # pytorch_lightning <= 1.7
+    from pytorch_lightning.utilities.types import _PATH
+except ImportError:  # pytorch_lightning >= 1.8
+    try:
+        from lightning_lite.utilities.types import _PATH
+    except ImportError:  # pytorch_lightning >= 1.9
+        from lightning_fabric.utilities.types import _PATH
 
 
 class DistAdamNativeMixedPrecisionPlugin(NativeMixedPrecisionPlugin):