 --enable_xformers_memory_efficient_attention
  --lr_scheduler="constant" --lr_warmup_steps=0 \
  --output_dir="sd-pokemon-model" \
  --push_to_hub
</hfoption>
<hfoption id="Flax">
Training with Flax can be faster on TPUs and GPUs thanks to @duongna211. Flax is more efficient on a TPU, but GPU performance is also great. Set the environment variables MODEL_NAME and dataset_name to the model and the dataset (either from the Hub or a local path). 