is only supported by PyTorch and is unavailable in the Flax training script. Add the --snr_gamma parameter and set it to the recommended value of 5.0:   Copied accelerate launch train_text_to_image.py \
  --snr_gamma=5.0 You can compare the loss surfaces for different snr_gamma values in this Weights and Biases report. For smaller datasets, the effects of Min-SNR may not be as obvious compared to larger datasets.  Training script The dataset pr