Run augment.py and out of memory after one epoch

[weihui98]

Hello,
Thanks for your great work!
When I run the augment.py, it is out of memory after the first epoch, is it the problem of the batch size?
The args are as follow:

09/11 03:17:31 AM |
09/11 03:17:31 AM | Parameters:
09/11 03:17:31 AM | AUX_WEIGHT=0.4
09/11 03:17:31 AM | BATCH_SIZE=48
09/11 03:17:31 AM | CUTOUT_LENGTH=16
09/11 03:17:31 AM | DATA_PATH=./data/
09/11 03:17:31 AM | DATASET=cifar10
09/11 03:17:31 AM | DROP_PATH_PROB=0.2
09/11 03:17:31 AM | EPOCHS=600
09/11 03:17:31 AM | GENOTYPE=Genotype(normal=[[('sep_conv_3x3', 1), ('sep_conv_3x3', 0)], [('sep_conv_3x3', 0), ('sep_conv_5x5', 1)],
[('sep_conv_3x3', 1), ('sep_conv_3x3', 2)], [('sep_conv_3x3', 0), ('sep_conv_3x3', 2)]], normal_concat=range(2, 6), reduce=[[('max_p
ool_3x3', 0), ('max_pool_3x3', 1)], [('max_pool_3x3', 0), ('skip_connect', 2)], [('max_pool_3x3', 0), ('skip_connect', 2)], [('max_po
ol_3x3', 0), ('skip_connect', 2)]], reduce_concat=range(2, 6))
09/11 03:17:31 AM | GPUS=[0]
09/11 03:17:31 AM | GRAD_CLIP=5.0
09/11 03:17:31 AM | INIT_CHANNELS=36
09/11 03:17:31 AM | LAYERS=20
09/11 03:17:31 AM | LR=0.025
09/11 03:17:31 AM | MOMENTUM=0.9
09/11 03:17:31 AM | NAME=cifar10
09/11 03:17:31 AM | PATH=augments/cifar10
09/11 03:17:31 AM | PRINT_FREQ=200
09/11 03:17:31 AM | SEED=2
09/11 03:17:31 AM | WEIGHT_DECAY=0.0003
09/11 03:17:31 AM | WORKERS=4
09/11 03:17:31 AM |
09/11 03:17:31 AM | Logger is set - training start