Skip to content

Commit 25ed51d

Browse files
pre-commit-ci[bot]jingxu10
authored andcommitted
[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
1 parent bc9da00 commit 25ed51d

File tree

1 file changed

+4
-5
lines changed

1 file changed

+4
-5
lines changed

src/lightning/fabric/strategies/deepspeed.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,11 +33,11 @@
3333
from lightning.fabric.strategies.registry import _StrategyRegistry
3434
from lightning.fabric.strategies.strategy import _Sharded
3535
from lightning.fabric.utilities.distributed import log
36+
from lightning.fabric.utilities.imports import _LIGHTNING_XPU_AVAILABLE
3637
from lightning.fabric.utilities.rank_zero import rank_zero_info, rank_zero_warn
3738
from lightning.fabric.utilities.seed import reset_seed
3839
from lightning.fabric.utilities.types import _PATH
3940

40-
from lightning.fabric.utilities.imports import _LIGHTNING_XPU_AVAILABLE
4141
if _LIGHTNING_XPU_AVAILABLE:
4242
from lightning_xpu.fabric import XPUAccelerator
4343

@@ -483,10 +483,9 @@ def load_checkpoint(
483483
optimzer_state_requested = bool(len([item for item in state.values() if isinstance(item, Optimizer)]))
484484

485485
torch.cuda.empty_cache()
486-
try:
486+
with suppress(AttributeError):
487487
torch.xpu.empty_cache()
488-
except AttributeError:
489-
pass
488+
490489
_, client_state = engine.load_checkpoint(
491490
path,
492491
tag="checkpoint",
@@ -584,7 +583,7 @@ def _initialize_engine(
584583
def _setup_distributed(self) -> None:
585584
if not isinstance(self.accelerator, CUDAAccelerator) and not isinstance(self.accelerator, XPUAccelerator):
586585
raise RuntimeError(
587-
f"The DeepSpeed strategy is only supported on CUDA/Intel(R) GPUs but"
586+
"The DeepSpeed strategy is only supported on CUDA/Intel(R) GPUs but"
588587
" `{self.accelerator.__class__.__name__}` is used."
589588
)
590589
reset_seed()

0 commit comments

Comments
 (0)