diff --git a/src/dvclive/fabric.py b/src/dvclive/fabric.py index 0e0903d..f8c7c8b 100644 --- a/src/dvclive/fabric.py +++ b/src/dvclive/fabric.py @@ -11,7 +11,7 @@ ) from lightning.fabric.utilities.rank_zero import rank_zero_only except ImportError: - from lightning_fabric.loggers.logger import Logger, rank_zero_experiment + from lightning_fabric.loggers.logger import Logger, rank_zero_experiment # type: ignore[assignment] from lightning_fabric.utilities.logger import ( _add_prefix, _convert_params, diff --git a/src/dvclive/lightning.py b/src/dvclive/lightning.py index dc15952..884ac16 100644 --- a/src/dvclive/lightning.py +++ b/src/dvclive/lightning.py @@ -13,7 +13,7 @@ from lightning.pytorch.utilities import rank_zero_only except ImportError: from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint # type: ignore[assignment] - from pytorch_lightning.loggers.logger import Logger + from pytorch_lightning.loggers.logger import Logger # type: ignore[assignment] from pytorch_lightning.utilities import rank_zero_only try: