[training] set rest of the blocks with requires_grad False. (#10607)

set rest of the blocks with requires_grad False.
huggingface · Jan 19, 2025 · 328e0d2 · 328e0d2
1 parent 23b467c
commit 328e0d2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/flux-control/train_control_flux.py b/examples/flux-control/train_control_flux.py
@@ -812,6 +812,8 @@ def main(args):
         for name, module in flux_transformer.named_modules():
             if "transformer_blocks" in name:
                 module.requires_grad_(True)
+            else:
+                module.requirs_grad_(False)
 
     def unwrap_model(model):
         model = accelerator.unwrap_model(model)