From 35358a2decad0c1e2972a8311ff49bfc14a326e1 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Wed, 8 May 2024 07:59:08 -1000 Subject: [PATCH] fix offload test (#7868) fix Co-authored-by: Dhruv Nair --- tests/pipelines/test_pipelines_common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 032fbb81ea..c4b7a3b918 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -1375,7 +1375,7 @@ class PipelineTesterMixin: output_without_offload = pipe(**inputs)[0] pipe.enable_sequential_cpu_offload() - assert pipe._execution_device.type == pipe._offload_device.type + assert pipe._execution_device.type == "cuda" inputs = self.get_dummy_inputs(generator_device) output_with_offload = pipe(**inputs)[0] @@ -1440,7 +1440,7 @@ class PipelineTesterMixin: output_without_offload = pipe(**inputs)[0] pipe.enable_model_cpu_offload() - assert pipe._execution_device.type == pipe._offload_device.type + assert pipe._execution_device.type == "cuda" inputs = self.get_dummy_inputs(generator_device) output_with_offload = pipe(**inputs)[0]