[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

src/lightning_fabric/strategies/deepspeed.py

@@ -19,7 +19,7 @@
 from contextlib import contextmanager
 from itertools import chain
 from pathlib import Path
-from typing import Any, Dict, Generator, Iterable, List, Mapping, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Any, Dict, Generator, List, Mapping, Optional, Tuple, TYPE_CHECKING, Union
 
 import torch
 from lightning_utilities.core.imports import RequirementCache
@@ -428,6 +428,7 @@ def load_checkpoint(
         torch.cuda.empty_cache()
 
         from deepspeed import DeepSpeedEngine
+
         modules = chain(*(module.modules() for module in state.values() if isinstance(module, Module)))
         engines = [engine for engine in modules if isinstance(engine, DeepSpeedEngine)]
         print(list(modules))
@@ -448,7 +449,10 @@ def load_checkpoint(
 
         optimzer_state_requested = bool(len([item for item in state.values() if isinstance(item, Optimizer)]))
         _, client_state = engine.load_checkpoint(
-            path, tag="checkpoint", load_optimizer_states=optimzer_state_requested, load_lr_scheduler_states=False,
+            path,
+            tag="checkpoint",
+            load_optimizer_states=optimzer_state_requested,
+            load_lr_scheduler_states=False,
             load_module_strict=True,  # TODO: make strict loading configurable
         )
         if client_state is None:

tests/tests_fabric/strategies/test_deepspeed_integration.py

@@ -346,9 +346,10 @@ def step(self, model, batch):
 
 
 def _assert_saved_model_is_equal(fabric, model, checkpoint_path):
-    """Convert the saved checkpoint to a single file with the model weights consolidated to easily verify the
-    full weights in float32 precision."""
+    """Convert the saved checkpoint to a single file with the model weights consolidated to easily verify the full
+    weights in float32 precision."""
     from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
+
     assert isinstance(fabric.strategy, DeepSpeedStrategy)
 
     # carry out the check only on rank 0