From 1aa3e050f475936a3a3051af23ca7f1b57debd8a Mon Sep 17 00:00:00 2001 From: Dave Berenbaum Date: Wed, 16 Aug 2023 05:25:12 -0400 Subject: [PATCH] fix broken pytorch-lightning import (#664) --- src/dvclive/lightning.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/dvclive/lightning.py b/src/dvclive/lightning.py index 5f0c1fff..bc8cb1e9 100644 --- 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