From 7b10e4ae65cc5830c581fba58638f5afb6e587cf Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 8 Aug 2025 13:34:29 +0530 Subject: [PATCH] [tests] device placement for non-denoiser components in group offloading LoRA tests (#12103) up --- tests/lora/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/lora/utils.py b/tests/lora/utils.py index 1d0c83751d..f09f0d8ecb 100644 --- a/tests/lora/utils.py +++ b/tests/lora/utils.py @@ -2400,7 +2400,6 @@ class PeftLoraLoaderMixinTests: components, _, _ = self.get_dummy_components(self.scheduler_classes[0]) pipe = self.pipeline_class(**components) - pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) denoiser = pipe.transformer if self.unet_kwargs is None else pipe.unet @@ -2416,6 +2415,10 @@ class PeftLoraLoaderMixinTests: num_blocks_per_group=1, use_stream=use_stream, ) + # Place other model-level components on `torch_device`. + for _, component in pipe.components.items(): + if isinstance(component, torch.nn.Module): + component.to(torch_device) group_offload_hook_1 = _get_top_level_group_offload_hook(denoiser) self.assertTrue(group_offload_hook_1 is not None) output_1 = pipe(**inputs, generator=torch.manual_seed(0))[0]