diff --git a/src/diffusers/hooks/group_offloading.py b/src/diffusers/hooks/group_offloading.py index 891ac28455af..30c0834a9e05 100644 --- a/src/diffusers/hooks/group_offloading.py +++ b/src/diffusers/hooks/group_offloading.py @@ -246,6 +246,8 @@ def _offload_to_memory(self): for group_module in self.modules: for param in group_module.parameters(): param.data = self.cpu_param_dict[param] + for buffer in group_module.buffers(): + buffer.data = self.cpu_param_dict[buffer] for param in self.parameters: param.data = self.cpu_param_dict[param] for buffer in self.buffers: