Update flux_train.py

kohya-ss · Aug 17, 2024 · 843e7e6 · 843e7e6
1 parent a5af5a0
commit 843e7e6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/flux_train.py b/flux_train.py
@@ -310,7 +310,7 @@ def train(args):
         logger.info(f"using {len(optimizers)} optimizers for fused optimizer groups")
 
     else:
-        _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize)
+        _, _, optimizer = train_util.get_optimizer(args, trainable_params=params_to_optimize, model=flux)
 
     # prepare dataloader
     # strategies are set here because they cannot be referenced in another process. Copy them with the dataset