File tree Expand file tree Collapse file tree 1 file changed +4
-5
lines changed
src/lightning/fabric/strategies Expand file tree Collapse file tree 1 file changed +4
-5
lines changed Original file line number Diff line number Diff line change 33
33
from lightning .fabric .strategies .registry import _StrategyRegistry
34
34
from lightning .fabric .strategies .strategy import _Sharded
35
35
from lightning .fabric .utilities .distributed import log
36
+ from lightning .fabric .utilities .imports import _LIGHTNING_XPU_AVAILABLE
36
37
from lightning .fabric .utilities .rank_zero import rank_zero_info , rank_zero_warn
37
38
from lightning .fabric .utilities .seed import reset_seed
38
39
from lightning .fabric .utilities .types import _PATH
39
40
40
- from lightning .fabric .utilities .imports import _LIGHTNING_XPU_AVAILABLE
41
41
if _LIGHTNING_XPU_AVAILABLE :
42
42
from lightning_xpu .fabric import XPUAccelerator
43
43
@@ -483,10 +483,9 @@ def load_checkpoint(
483
483
optimzer_state_requested = bool (len ([item for item in state .values () if isinstance (item , Optimizer )]))
484
484
485
485
torch .cuda .empty_cache ()
486
- try :
486
+ with suppress ( AttributeError ) :
487
487
torch .xpu .empty_cache ()
488
- except AttributeError :
489
- pass
488
+
490
489
_ , client_state = engine .load_checkpoint (
491
490
path ,
492
491
tag = "checkpoint" ,
@@ -584,7 +583,7 @@ def _initialize_engine(
584
583
def _setup_distributed (self ) -> None :
585
584
if not isinstance (self .accelerator , CUDAAccelerator ) and not isinstance (self .accelerator , XPUAccelerator ):
586
585
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"
588
587
" `{self.accelerator.__class__.__name__}` is used."
589
588
)
590
589
reset_seed ()
You can’t perform that action at this time.
0 commit comments