AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-29 07:22:12 +03:00
Code Activity

[Tests] Fix cpu offload test (#5626)

Browse Source 
* fix more

* fix more
This commit is contained in:
Patrick von Platen
2023-11-02 12:49:58 +01:00
committed by GitHub
parent b81f709fb6
commit 9723f8a557
1 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/pipelines/test_pipelines_common.py
View File

@@ -742,14 +742,14 @@ class PipelineTesterMixin:
max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max()
self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results")
offloaded_modules = [
v
for k, v in pipe.components.items()
if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload
]
self.assertTrue(
all(
v.device == "cpu"
for k, v in pipe.components.values()
if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload
),
"CPU offloading should leave all pipeline components on the CPU after inference",
)
all(v.device.type == "cpu" for v in offloaded_modules)
), f"Not offloaded: {[v for v in offloaded_modules if v.device.type != 'cpu']}"
@unittest.skipIf(
torch_device != "cuda" or not is_xformers_available(),