update

Author: justusschock

_notebooks

@@ -346,7 +346,7 @@ def model_sharded_context(self) -> Generator[None, None, None]:
     @property
     @override
     def distributed_sampler_kwargs(self) -> dict[str, int]:
-        return self.deepspeed_strategy_impl.distributed_sampler_kwargs()
+        return self.deepspeed_strategy_impl.distributed_sampler_kwargs
 
     @override
     def setup_optimizers(self, trainer: "pl.Trainer") -> None:

src/lightning/pytorch/strategies/deepspeed.py

@@ -16,6 +16,7 @@
 from unittest.mock import MagicMock, Mock
 
 import pytest
+import pytorch_lightning_enterprise.loggers  # noqa: F401
 from pytorch_lightning_enterprise.utils.imports import _MLFLOW_AVAILABLE
 
 from lightning.pytorch import Trainer