AttributeError: объект «SparseCategoricalCrossentropy» не имеет атрибута «_id»

Я пытался воссоздать простую DNN, используя только базовый слой Keras и пишу все с нуля. Вроде все работает нормально, но во время цикла обучения я получаю эту ошибку: AttributeError: 'SparseCategoricalCrossentropy' object has no attribute '_id'

Я пытался изменить функцию потерь на CategoricalCrossentropy и SparseCategoricalCrossentropy (с from_logits True или False), но ошибка всегда появляется.

Вот код:

import numpy as np
import tensorflow as tf
from tensorflow import keras

from utils import plot_image, plot_mnist_results, plot_value_array


class Flatten(keras.layers.Layer):
    def __init__(self):
        super(Flatten, self).__init__()

    def build(self, input_shape):
        self.output_size = np.prod(input_shape)

    def call(self, X):
        return tf.reshape(X, shape=(-1, self.output_size))


class Dense(keras.layers.Layer):
    def __init__(self, units, activation):
        super(Dense, self).__init__()
        self.units = units
        self.activation = activation

    def build(self, input_shape):
        self.kernel = self.add_weight(
            name='kernel',
            dtype=tf.float64,
            initializer='glorot_normal',
            trainable=True,
            shape=(input_shape[-1], self.units)
        )
        self.bias = self.add_weight(
            name='bias',
            dtype=tf.float64,
            initializer=keras.initializers.Constant(0.1),
            trainable=True,
            shape=(1, self.units)
        )

    def call(self, X):
        return self.activation(tf.matmul(X, self.kernel) + self.bias)


class DNN(keras.models.Model):
    def __init__(self, units, activation):
        super(DNN, self).__init__()
        self.units = units
        self.activation = activation

    def build(self, input_shape):
        self.flatten = Flatten()
        self.hidden_layer = Dense(self.units, tf.nn.relu)
        self.output_layer = Dense(10, tf.nn.softmax)

    def call(self, X):
        print(self.hidden_layer(self.flatten(X)).shape)
        print(self.output_layer(self.hidden_layer(self.flatten(X))).shape)
        return self.output_layer(self.hidden_layer(self.flatten(X)))


# @tf.function
def train(model, loss, opt, X, y):
    with tf.GradientTape() as tape:
        gradients = tape.gradient(loss(model(X), y), model.trainable_variables)
        gradient_variables = zip(gradients, model.trainable_variables)
        opt.apply_gradients(gradient_variables)


mnist = keras.datasets.mnist

(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

train_images = train_images / 255.0
test_images = test_images / 255.0

model = DNN(units=128, activation=tf.nn.relu)
opt = tf.optimizers.Adam(learning_rate=1e-3)

for epoch in range(3):

    for step in range(train_labels.shape[0]):
        loss = keras.losses.SparseCategoricalCrossentropy
        train(model, loss, opt, train_images[step, :, :], train_labels[step])

    train_loss = loss(model(train_images), train_labels)

    template = 'Epoch {}, Train loss: {:.5f}'
    print(template.format(epoch + 1, train_loss.numpy()))

Я ожидаю, что модель будет успешно обучаться, но, похоже, это не так. Что я делаю не так?