File tree Expand file tree Collapse file tree 1 file changed +3
-3
lines changed Expand file tree Collapse file tree 1 file changed +3
-3
lines changed Original file line number Diff line number Diff 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
You can’t perform that action at this time.
0 commit comments