From 328e0d20a7b996f9bdb04180457eb08c1b42a76e Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sun, 19 Jan 2025 19:34:53 +0530 Subject: [PATCH] [training] set rest of the blocks with `requires_grad` False. (#10607) set rest of the blocks with requires_grad False. --- examples/flux-control/train_control_flux.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/flux-control/train_control_flux.py b/examples/flux-control/train_control_flux.py index 7d0e28069054..4449811ab747 100644 --- 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)