Unverified Commit 9723f8a5 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

[Tests] Fix cpu offload test (#5626)

* fix more

* fix more
parent b81f709f
......@@ -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")
self.assertTrue(
all(
v.device == "cpu"
for k, v in pipe.components.values()
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
),
"CPU offloading should leave all pipeline components on the CPU after inference",
)
]
self.assertTrue(
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(),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment