Skip to content

Commit ecf2321

Browse files
authored
Teardown from lightning-xpu (from #PR- 3)
From Lightning-AI#3
1 parent 0040f01 commit ecf2321

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

src/lightning/fabric/strategies/deepspeed.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -484,8 +484,9 @@ def load_checkpoint(
484484

485485
torch.cuda.empty_cache()
486486
with suppress(AttributeError):
487-
torch.xpu.empty_cache()
488-
487+
if _LIGHTNING_XPU_AVAILABLE:
488+
XPUAccelerator.teardown()
489+
489490
_, client_state = engine.load_checkpoint(
490491
path,
491492
tag="checkpoint",

0 commit comments

Comments
 (0)