Malloc error while training on birds dataset

[parthrangarajan]

So I have followed all the steps necessary till the training of the StackGAN on the birds dataset but I am receiving a memory allocation error while training. I will enclose the Colab link I am using as well as the detailed error statement for your convenience. Any help would be deeply appreciated. Thank you!
Regards,
Parth Rangarajan.
Here is the collaboratory link: https://colab.research.google.com/drive/1ASXNdAWI54x8zXZROHcYhkhEX8ciK8-H?usp=sharing

Here is the error:
Using config: {'CONFIG_NAME': '3stages', 'CUDA': True, 'DATASET_NAME': 'birds', 'DATA_DIR': '../data/birds', 'EMBEDDING_TYPE': 'cnn-rnn', 'GAN': {'B_CONDITION': True, 'DF_DIM': 64, 'EMBEDDING_DIM': 128, 'GF_DIM': 64, 'NETWORK_TYPE': 'default', 'R_NUM': 2, 'Z_DIM': 100}, 'GPU_ID': '0', 'TEST': {'B_EXAMPLE': True, 'SAMPLE_NUM': 30000}, 'TEXT': {'DIMENSION': 1024}, 'TRAIN': {'BATCH_SIZE': 24, 'COEFF': {'COLOR_LOSS': 0.0, 'KL': 2.0, 'UNCOND_LOSS': 1.0}, 'DISCRIMINATOR_LR': 0.0002, 'FLAG': True, 'GENERATOR_LR': 0.0002, 'MAX_EPOCH': 600, 'NET_D': '', 'NET_G': '', 'SNAPSHOT_INTERVAL': 2000, 'VIS_COUNT': 64}, 'TREE': {'BASE_SIZE': 64, 'BRANCH_NUM': 3}, 'WORKERS': 4} /usr/local/lib/python3.7/dist-packages/torchvision/transforms/transforms.py:310: UserWarning: The use of the transforms.Scale transform is deprecated, please use transforms.Resize instead. warnings.warn("The use of the transforms.Scale transform is deprecated, " + Total filenames: 11788 001.Black_footed_Albatross/Black_Footed_Albatross_0046_18.jpg Load filenames from: ../data/birds/train/filenames.pickle (8855) tcmalloc: large alloc 65850621952 bytes == 0x5631f8284000 @ 0x7fd36b19a001 0x7fd274fd754f 0x7fd275027b58 0x7fd27502ae83 0x7fd27502b07b 0x7fd2750cc761 0x5631778c24b0 0x5631778c2240 0x5631779360f3 0x5631778c3afa 0x563177931c0d 0x5631779309ee 0x5631778c448c 0x563177905159 0x5631779020a4 0x5631778c2d49 0x56317793694f 0x5631779309ee 0x5631779306f3 0x5631779fa4c2 0x5631779fa83d 0x5631779fa6e6 0x5631779d2163 0x5631779d1e0c 0x7fd369f82bf7 0x5631779d1cea ^C