Unverified Commit e2bbaa4f authored by statelesshz's avatar statelesshz Committed by GitHub
Browse files

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
parent 1e853e24
...@@ -1127,7 +1127,9 @@ class DiffusionPipeline(ConfigMixin): ...@@ -1127,7 +1127,9 @@ class DiffusionPipeline(ConfigMixin):
if self.device.type != "cpu": if self.device.type != "cpu":
self.to("cpu", silence_dtype_warnings=True) self.to("cpu", silence_dtype_warnings=True)
torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist) device_mod = getattr(torch, self.device.type, None)
if hasattr(device_mod, "empty_cache") and device_mod.is_available():
device_mod.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
for name, model in self.components.items(): for name, model in self.components.items():
if not isinstance(model, torch.nn.Module): if not isinstance(model, torch.nn.Module):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment