如何在tensorflow 2.0中训练以数组为标签的神经网络，即[1,0]

import tensorflow as tf
from tensorflow import keras

train_data = [...]  # 输入数据
train_labels = [[1, 0], [0, 1], ...]  # 数组标签

model = keras.Sequential([
    keras.layers.Dense(64, activation='relu', input_shape=(input_shape,)),
    keras.layers.Dense(2, activation='softmax')
])

model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

model.fit(train_data, train_labels, epochs=10, batch_size=32)