Error when I'm trying to train a Flux lora with train_dreambooth_lora_flux_advanced

[Mnwa]

Describe the bug

Hi! I'm trying to train my lora model with train_dreambooth_lora_flux_advanced script.

When I'm trying to train my model with prior preservation tag I give an error.

How can I fix it?

Reproduction

accelerate launch train_dreambooth_lora_flux_advanced.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --dataset_name="./ds5" \
  --instance_prompt="1boy, 1girl" \
  --validation_prompt="1boy, 1girl" \
  --class_prompt="1boy, 1girl" \
  --num_class_images=200 \
  --with_prior_preservation \
  --class_data_dir="./cdi" \
  --output_dir="crtr-SDXL-LoRA" \
  --caption_column="text" \
  --mixed_precision="bf16" \
  --prior_generation_precision="bf16" \
  --resolution=1024 \
  --train_batch_size=8 \
  --repeats=1 \
  --gradient_accumulation_steps=8 \
  --gradient_checkpointing \
  --learning_rate=1.0 \
  --optimizer="prodigy"\
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --rank=64 \
  --num_train_epochs=200 \
  --validation_epochs=100 \
  --center_crop \
  --adam_beta2=0.99 \
  --adam_weight_decay=0.01 \
  --allow_tf32

Logs

Traceback (most recent call last):
  File "/workspace/train_dreambooth_lora_flux_advanced.py", line 2423, in <module>
    main(args)
  File "/workspace/train_dreambooth_lora_flux_advanced.py", line 2213, in main
    (weighting.float() * (model_pred_prior.float() - target_prior.float()) ** 2).reshape(
     ~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (16) must match the size of tensor b (8) at non-singleton dimension 0

System Info

Diffusers 0.33
CUDA 12.9
Torch 2.7

Docker image
nvcr.io/nvidia/pytorch:25.04-py3

Who can help?

@sayakpaul

[Mnwa]

Same error when I'm training a SDXL model with --do_edm_style_training flag

[sayakpaul]

Cc: @linoytsaban