Skip to content

Commit 35358a2

Browse files
yiyixuxuDN6
andauthored
fix offload test (#7868)
fix Co-authored-by: Dhruv Nair <[email protected]>
1 parent 818f760 commit 35358a2

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

tests/pipelines/test_pipelines_common.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1375,7 +1375,7 @@ def test_sequential_cpu_offload_forward_pass(self, expected_max_diff=1e-4):
13751375
output_without_offload = pipe(**inputs)[0]
13761376

13771377
pipe.enable_sequential_cpu_offload()
1378-
assert pipe._execution_device.type == pipe._offload_device.type
1378+
assert pipe._execution_device.type == "cuda"
13791379

13801380
inputs = self.get_dummy_inputs(generator_device)
13811381
output_with_offload = pipe(**inputs)[0]
@@ -1440,7 +1440,7 @@ def test_model_cpu_offload_forward_pass(self, expected_max_diff=2e-4):
14401440
output_without_offload = pipe(**inputs)[0]
14411441

14421442
pipe.enable_model_cpu_offload()
1443-
assert pipe._execution_device.type == pipe._offload_device.type
1443+
assert pipe._execution_device.type == "cuda"
14441444

14451445
inputs = self.get_dummy_inputs(generator_device)
14461446
output_with_offload = pipe(**inputs)[0]

0 commit comments

Comments
 (0)