diff --git a/examples/z_image/model_training/full/Z-Image-Omni-Base.sh b/examples/z_image/model_training/full/Z-Image-Omni-Base.sh
index cc74b2a..4ed3c18 100644
--- a/examples/z_image/model_training/full/Z-Image-Omni-Base.sh
+++ b/examples/z_image/model_training/full/Z-Image-Omni-Base.sh
@@ -12,6 +12,7 @@ accelerate launch --config_file examples/z_image/model_training/full/accelerate_
   --output_path "./models/train/Z-Image-Omni-Base_full" \
   --trainable_models "dit" \
   --use_gradient_checkpointing \
+  --find_unused_parameters \
   --dataset_num_workers 8
 
 # Image(s) to image training
@@ -29,4 +30,5 @@ accelerate launch --config_file examples/z_image/model_training/full/accelerate_
 #   --output_path "./models/train/Z-Image-Omni-Base_full_edit" \
 #   --trainable_models "dit" \
 #   --use_gradient_checkpointing \
+#   --find_unused_parameters \
 #   --dataset_num_workers 8
diff --git a/examples/z_image/model_training/lora/Z-Image-Omni-Base.sh b/examples/z_image/model_training/lora/Z-Image-Omni-Base.sh
index ef4d524..d8f522a 100644
--- a/examples/z_image/model_training/lora/Z-Image-Omni-Base.sh
+++ b/examples/z_image/model_training/lora/Z-Image-Omni-Base.sh
@@ -13,6 +13,7 @@ accelerate launch examples/z_image/model_training/train.py \
   --lora_target_modules "to_q,to_k,to_v,to_out.0,w1,w2,w3" \
   --lora_rank 32 \
   --use_gradient_checkpointing \
+  --find_unused_parameters \
   --dataset_num_workers 8
 
 # Image(s) to image training
@@ -32,4 +33,5 @@ accelerate launch examples/z_image/model_training/train.py \
 #   --lora_target_modules "to_q,to_k,to_v,to_out.0,w1,w2,w3" \
 #   --lora_rank 32 \
 #   --use_gradient_checkpointing \
+#   --find_unused_parameters \
 #   --dataset_num_workers 8