ValueError: Found array with dim 3. Estimator expected <= 2. on large number of epochs code example

Example: ValueError: Found array with dim 3. Estimator expected <= 2.

nsamples, nx, ny = train_dataset.shape
d2_train_dataset = train_dataset.reshape((nsamples,nx*ny))