Skip to content

GitLab

Commit b705afb9 authored by muyangli's avatar muyangli
Browse files

uses model cpu offload instead of per-layer offload to speedup test ci

parent e8de2628
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
tests/flux/utils.py
View file @ b705afb9
...@@ -191,7 +191,7 @@ def run_test( ...@@ -191,7 +191,7 @@ def run_test(
if gpu_memory > 36 * 1024: if gpu_memory > 36 * 1024:
pipeline = pipeline.to("cuda") pipeline = pipeline.to("cuda")
else: else:
pipeline.enable_sequential_cpu_offload() pipeline.enable_model_cpu_offload()
if len(lora_names) > 0: if len(lora_names) > 0:
for i, (lora_name, lora_strength) in enumerate(zip(lora_names, lora_strengths)): for i, (lora_name, lora_strength) in enumerate(zip(lora_names, lora_strengths)):
......
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