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

@sayakpaul

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

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

@linoytsaban

Cc: @linoytsaban

If I'm splitting weighting

if args.with_prior_preservation:
    weighting_instance, weighting_prior = torch.chunk(weighting, 2, dim=0)
else:
    weighting_instance = weighting

And using divided weights

if args.with_prior_preservation:
    model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
    target, target_prior = torch.chunk(target, 2, dim=0)

    prior_loss = torch.mean(
        (weighting_prior.float() * (model_pred_prior.float() - target_prior.float()) ** 2).reshape(
            target_prior.shape[0], -1
        ),
        1,
    )
    prior_loss = prior_loss.mean()

loss = torch.mean(
    (weighting_instance.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
    1,
)
loss = loss.mean()

Then it works, but I'm not sure that is correct

Hi! Can you help about my solution? Is it correct?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

Navigation Menu

Repository navigation

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

Describe the bug

Reproduction

Logs

System Info

Who can help?

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

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

Description

Describe the bug

Reproduction

Logs

System Info

Who can help?

Activity

Mnwa commented on May 4, 2025

sayakpaul commented on May 5, 2025

Mnwa commented on May 6, 2025

Mnwa commented on Jul 6, 2025

github-actions commented on Jan 9, 2026

github-actions commented on Feb 3, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions