Skip to content

Commit e2bbaa4

Browse files
authored
make enable_sequential_cpu_offload more generic for third-party devices (#4191)
* make enable_sequential_cpu_offload more generic for third-party devices * make style
1 parent 1e853e2 commit e2bbaa4

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

src/diffusers/pipelines/pipeline_utils.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1127,7 +1127,9 @@ def enable_sequential_cpu_offload(self, gpu_id: int = 0, device: Union[torch.dev
11271127

11281128
if self.device.type != "cpu":
11291129
self.to("cpu", silence_dtype_warnings=True)
1130-
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
1130+
device_mod = getattr(torch, self.device.type, None)
1131+
if hasattr(device_mod, "empty_cache") and device_mod.is_available():
1132+
device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
11311133

11321134
for name, model in self.components.items():
11331135
if not isinstance(model, torch.nn.Module):

0 commit comments

Comments
 (0)