From 7c88e5fe21ab74fe9247ebef32ee903796efedae Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 22 Apr 2026 11:41:37 -0300 Subject: [PATCH] [tests] fix group offloading with disk tests (#13491) fix group offloading with disk tests --- tests/models/testing_utils/memory.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/models/testing_utils/memory.py b/tests/models/testing_utils/memory.py index 68480bddd3..8731c64485 100644 --- a/tests/models/testing_utils/memory.py +++ b/tests/models/testing_utils/memory.py @@ -22,10 +22,10 @@ import pytest import torch from accelerate.utils.modeling import compute_module_sizes -from diffusers.utils.testing_utils import _check_safetensors_serialization from diffusers.utils.torch_utils import get_torch_cuda_device_capability from ...testing_utils import ( + _check_safetensors_serialization, assert_tensors_close, backend_empty_cache, backend_max_memory_allocated, @@ -361,6 +361,9 @@ class GroupOffloadTesterMixin: offload_to_disk_path=tmpdir, offload_type=offload_type, num_blocks_per_group=num_blocks_per_group, + block_modules=model._group_offload_block_modules + if hasattr(model, "_group_offload_block_modules") + else None, ) if not is_correct: if extra_files: