How to use sparse_categorical_cross_entropy in an autoencoder?

I am trying to train an autoencoder and due to memory issues I decided that I would stop using one-hot-encoding, instead tryout the SparseCategoricalCrossentropy(). Unfortunetely after trying to train the model I get: tensorflow.python.framework.errors_impl.InvalidArgumentError: logits and labels must have the same first dimension, got logits shape [512,32] and labels shape [2560]

from keras.metrics import SparseCategoricalAccuracy
from tensorflow.keras.layers import *
from tensorflow.keras.models import *
from tensorflow.keras.losses import *

import numpy as np
from tensorflow.python.keras.optimizer_v1 import RMSprop

BATCH_SIZE = 128
EPOCHS = 10
numpy_dataset = np.create
def create_autoencoder(sequence_length, latent_dim):
    # Define encoder input shape
    encoder_input = Input(shape=sequence_length, batch_size=BATCH_SIZE)
    latent = Input(shape=latent_dim)
    encoded = Dense(latent_dim, activation='tanh')(encoder_input)
    decoded = Dense(sequence_length, activation='sigmoid')(latent)
    encoder = Model(inputs=encoder_input, outputs=encoded, name="encoder")
    decoder = Model(inputs=latent, outputs=decoded, name="decoder")
    autoencoder = Model(inputs=encoder_input, outputs=decoder(encoded), name="autoencoder")
    return autoencoder


def data_generator(numpy_dataset, steps_per_epoch):
    # Computes the number of elements that should be removed so that the training data can be split
    # equally between batches
    # elements2remove = len(train_notes) - steps_per_epoch * BATCH_SIZE

    # steps_per_epoch =
    for i in range(EPOCHS):
        for batch in np.array_split(numpy_dataset, steps_per_epoch):
            yield batch, batch


def train(model, numpy_dataset):
    steps_per_epoch = len(numpy_dataset) // BATCH_SIZE
    model.compile(loss=SparseCategoricalCrossentropy(),
                  optimizer=RMSprop(learning_rate=0.01),
                  metrics=[SparseCategoricalAccuracy()])

    model.fit(x=data_generator(numpy_dataset=numpy_dataset, steps_per_epoch=steps_per_epoch),
              epochs=EPOCHS,
              batch_size=BATCH_SIZE,
              steps_per_epoch=steps_per_epoch)

autoencoder = create_autoencoder(32, 2)
train(autoencoder, numpy_dataset)

I simplified the code, to make it clearer. So lets say that earlier in the code I create the numpy_dataset, and it is of shape=(2560, 32) and dtype=int32. How can I fix this issue?