我正在使用时尚的MNIST数据集，我想使用由我自己创建的自定义图层来构建人工神经网络：

这是我的预处理阶段：

#ANN with fashion mnist data_set
import tensorflow as tf
import numpy as np
from tensorflow.keras.datasets import fashion_mnist
from tensorflow.keras.layers import Layer
from tensorflow.keras import Model

#Loading dataset and preprocessing:
(X_train,Y_train),(X_test,Y_test) = fashion_mnist.load_data()

unique_labels = set(Y_train)

X_train = X_train/255 #because max value of pixel is 255
X_test = X_test/255 #because max value of pixel is 255

X_train = X_train.reshape(-1,28*28)#flatten image from 28*28 to N*(28*28),second dimension - all the pixels of image
X_test = X_test.reshape(-1,28*28)

Y_train = Y_train.reshape(Y_train.shape[0],1)
Y_test = Y_test.reshape(Y_test.shape[0],1)

dataset_train = tf.data.Dataset.from_tensor_slices((X_train,Y_train))
dataset_train = dataset.shuffle(buffer_size=1024).batch(64)

dataset_test = tf.data.Dataset.from_tensor_slices((X_test,Y_test))
dataset_test = dataset.shuffle(buffer_size=1024).batch(64)

在这里您可以看到，我已经对数据进行了预处理，以使其具有二维形状（更适合ANN）。 我知道，可以将cnn用于以下任务，但我的目标是与ANN一起练习。

最后，我得到以下错误：

InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-104-24b493b0dcc5> in <module>()
    102 
    103 for epoch in range(EPOCHS):
--> 104   for X,y in dataset_train:
    105     training(X,y)
    106 

4 frames
/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value,from_value)

InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [256,2] and element 36 had shape [32,2]. [Op:IteratorGetNextSync]

但是，如果我不使用batch(64)，则一切正常。 我该如何解决这个问题？