Skip to content

GitLab

Unverified Commit f1a32203 authored by Anton Lozhkov's avatar Anton Lozhkov Committed by GitHub
Browse files

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

parent 6f150263
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
tests/pipelines/unclip/test_unclip.py
View file @ f1a32203
......@@ -286,7 +286,7 @@ class UnCLIPPipelineIntegrationTests(unittest.TestCase):
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 @@ class UnCLIPPipelineIntegrationTests(unittest.TestCase):
)
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
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment