diff --git a/tests/pipelines/unclip/test_unclip.py b/tests/pipelines/unclip/test_unclip.py index aa5f63b72ced..c1f67e557fd9 100644 --- a/tests/pipelines/unclip/test_unclip.py +++ b/tests/pipelines/unclip/test_unclip.py @@ -286,7 +286,7 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self): torch.cuda.reset_max_memory_allocated() torch.cuda.reset_peak_memory_stats() - pipe = UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha") + pipe = UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha", torch_dtype=torch.float16) pipe = pipe.to(torch_device) pipe.set_progress_bar_config(disable=None) pipe.enable_attention_slicing() @@ -304,5 +304,5 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self): ) mem_bytes = torch.cuda.max_memory_allocated() - # make sure that less than 1.5 GB is allocated - assert mem_bytes < 1.5 * 10**9 + # make sure that less than 7 GB is allocated + assert mem_bytes < 7 * 10**9