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

b705afb9 · muyangli · e8de2628 · b705afb9
Commit b705afb9 authored Apr 12, 2025 by muyangli
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/flux/utils.py tests/flux/utils.py +1 -1

No files found.
--- a/tests/flux/utils.py
+++ b/tests/flux/utils.py
@@ -191,7 +191,7 @@ def run_test(
        if gpu_memory > 36 * 1024:
            pipeline = pipeline.to("cuda")
        else:
-            pipeline.enable_sequential_cpu_offload()
+            pipeline.enable_model_cpu_offload()

        if len(lora_names) > 0:
            for i, (lora_name, lora_strength) in enumerate(zip(lora_names, lora_strengths)):