From e2bbaa4f540a833a84b5644eee29577c3cba4287 Mon Sep 17 00:00:00 2001 From: statelesshz <371503430@qq.com> Date: Fri, 21 Jul 2023 19:45:09 +0800 Subject: [PATCH] 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 --- src/diffusers/pipelines/pipeline_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index ad52c6ac1c..3d827596d5 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -1127,7 +1127,9 @@ class DiffusionPipeline(ConfigMixin): if self.device.type != "cpu": 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(): if not isinstance(model, torch.nn.Module):