using Flux: onehotbatch

x_vec = [float.(reshape(x_train[:, :, :, i], :)) for i in 1:1000];
X = hcat(x_vec...);
Y = onehotbatch(y_train_fine[1:1000], 0:99);