BATCH_SIZE = 128
NUM_EPOCHS = 10