train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))

BATCH_SIZE = 64
train_dataset = train_dataset.shuffle(buffer_size=10000).batch(BATCH_SIZE)
test_dataset = test_dataset.batch(BATCH_SIZE)

        __
        __