Merge pull request kohya-ss#1322 from aria1th/patch-1

Accelerate: fix get_trainable_params in controlnet-llite training
feffy380 · May 19, 2024 · 0640f01 · 0640f01
2 parents 2f19175 + 793aeb9
commit 0640f01
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sdxl_train_control_net_lllite.py b/sdxl_train_control_net_lllite.py
@@ -477,7 +477,7 @@ def remove_model(old_ckpt_name):
 
                 accelerator.backward(loss)
                 if accelerator.sync_gradients and args.max_grad_norm != 0.0:
-                    params_to_clip = unet.get_trainable_params()
+                    params_to_clip = accelerator.unwrap_model(unet).get_trainable_params()
                     accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
 
                 optimizer.step()