fix broken pytorch-lightning import (#664)

iterative · Aug 16, 2023 · 1aa3e05 · 1aa3e05
1 parent a2cae20
commit 1aa3e05
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/dvclive/lightning.py b/src/dvclive/lightning.py
@@ -22,7 +22,11 @@
     from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
     from pytorch_lightning.loggers.logger import Logger, rank_zero_experiment
     from pytorch_lightning.utilities import rank_zero_only
-    from pytorch_lightning.utilities.logger import _scan_checkpoints
+
+    try:
+        from pytorch_lightning.utilities.logger import _scan_checkpoints
+    except ImportError:
+        from pytorch_lightning.loggers.utilities import _scan_checkpoints
 from torch import is_tensor
 
 from dvclive import Live