AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

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

Browse Source
This commit is contained in:
Anton Lozhkov
2022-12-19 18:25:08 +01:00
committed by GitHub
parent 6f15026330
commit f1a32203aa
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/pipelines/unclip/test_unclip.py
View File

@@ -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