using Flux train_batches = Flux.DataLoader((X_train_aug, y_train_one_hot), batchsize=BATCH_SIZE, shuffle=false)