How to create a Tensorflow dataset for image sequence prediction or next frame prediction from PNGs?

I'm trying to make a model for image sequence prediction using Tensorflow on Google Colab. Basically, the model should predict the next frames for a given image sequence.

But while fitting the model to the training data I get the following error:

ValueError: Input 0 of layer conv_lst_m2d_1 is incompatible with the layer: expected ndim=5, found ndim=4. Full shape received: (None, 900, 900, 3)

I followed this tutorial for the next frame prediction

https://keras.io/examples/vision/conv_lstm/

and tried to use this code for loading my images

https://www.tensorflow.org/tutorials/load_data/images

My images are PNGs (900x900px, rgb) and are stored like this, the filename is a unix-timestamp in miliseconds:

raw_images/
    1626008400000.png
    1626008700000.png
    1626009000000.png
    1626009300000.png
    1626009600000.png
    1626009900000.png
    ...

My code (minimal working version):

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
IMAGE_FOLDER = "raw_images/"

# Load the dataset
dataset = keras.preprocessing.image_dataset_from_directory(
    IMAGE_FOLDER,
    None,
    None,
    None,
    "rgb",
    32,
    (900, 900),
    False,
    None,
    None,
    None,
    "bilinear",
    False,
    True
)

dataset = dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)


model = tf.keras.Sequential([
  tf.keras.layers.experimental.preprocessing.Rescaling(1./255),
  tf.keras.layers.ConvLSTM2D(
    filters=64,
    kernel_size=(5, 5),
    padding="same",
    return_sequences=True,
    activation="relu",
),
  tf.keras.layers.BatchNormalization(),
  tf.keras.layers.ConvLSTM2D(
    filters=64,
    kernel_size=(3, 3),
    padding="same",
    return_sequences=True,
    activation="relu",
),
  tf.keras.layers.BatchNormalization(),
  tf.keras.layers.ConvLSTM2D(
    filters=64,
    kernel_size=(1, 1),
    padding="same",
    return_sequences=True,
    activation="relu",
),
  tf.keras.layers.Conv3D(
    filters=1, kernel_size=(3, 3, 3), activation="sigmoid", padding="same"
)
])


model.compile(
    loss=keras.losses.binary_crossentropy, 
    optimizer=keras.optimizers.Adam(),
)


# Define some callbacks to improve training.
early_stopping = keras.callbacks.EarlyStopping(monitor="val_loss", patience=10)
reduce_lr = keras.callbacks.ReduceLROnPlateau(monitor="val_loss", patience=5)

# Define modifiable training hyperparameters.
epochs = 20
batch_size = 32

# Fit the model to the training data.
model.fit(
    dataset,
    None,
    batch_size=batch_size,
    epochs=epochs
)

model.save("model")

I suspect I somehow need to add the timestamps to the dataset, but I can't find any way to do that.