mple images you provide. If you’re training on a GPU with limited vRAM, you should try enabling the gradient_checkpointing and mixed_precision parameters in the training command. You can also reduce your memory footprint by using memory-efficient attention with xFormers. JAX/Flax training is also supported for efficient training on TPUs and GPUs, but it doesn’t support gradient checkpointing or xFormers. With the same configuration and setup as