Skip to content

Commit f1a3220

Browse files
authored
[Tests] Fix UnCLIP cpu offload tests (huggingface#1769)
1 parent 6f15026 commit f1a3220

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

tests/pipelines/unclip/test_unclip.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -286,7 +286,7 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
286286
torch.cuda.reset_max_memory_allocated()
287287
torch.cuda.reset_peak_memory_stats()
288288

289-
pipe = UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha")
289+
pipe = UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha", torch_dtype=torch.float16)
290290
pipe = pipe.to(torch_device)
291291
pipe.set_progress_bar_config(disable=None)
292292
pipe.enable_attention_slicing()
@@ -304,5 +304,5 @@ def test_stable_diffusion_pipeline_with_sequential_cpu_offloading(self):
304304
)
305305

306306
mem_bytes = torch.cuda.max_memory_allocated()
307-
# make sure that less than 1.5 GB is allocated
308-
assert mem_bytes < 1.5 * 10**9
307+
# make sure that less than 7 GB is allocated
308+
assert mem_bytes < 7 * 10**9

0 commit comments

Comments
 (0)