BATCH_SIZE = 32
NUM_EPOCHS = 2