[Tests] Fix UnCLIP cpu offload tests (huggingface#1769)

0xdigiscore · Dec 19, 2022 · f1a3220 · f1a3220
1 parent 6f15026
commit f1a3220
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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