text_to_image multi-gpu not working

We are training text_to_image on Google cloud platform, the jupyterlab instance has 2 GPUs (NVIDIA Tesla P100) with a total memory of 32GB (16GB each). I tried using accelerate for training the text_to_image model for multi_gpu support. But still getting out of memory error. Even with 32GB, I don't understand why its only taking 16GB memory

Command used: `accelerate launch --multi_gpu train_text_to_image.py --pretrained_model_name_or_path=$MODEL_NAME --train_data_dir=$DATASET_DIR --image_column="image" --caption_column="text" --output_dir=$OUTPUT_DIR --train_batch_size=2 --resolution=512 --gradient_accumulation_steps=5 --num_train_epochs=1000 --learning_rate=1e-06 --gradient_checkpointing --enable_xformers_memory_efficient_attention`

`rank1]: torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 114.00 MiB. GPU has a total capacity of 15.89 GiB of which 89.12 MiB is free. Including non-PyTorch memory, this process has 15.80 GiB memory in use. Of the allocated memory 15.35 GiB is allocated by PyTorch, and 71.31 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)`


Any help will be much appreciated. Thanks.

Start with train_text_to_image.py and the supplied accelerate launch command, then inspect how multi-GPU training and batch allocation are initialized. Reproduce the command on two 16 GB P100 GPUs and determine whether both GPUs are used without the reported out-of-memory error.

索引モデルが issue の本文から書いたものです。