最后我这样解决了它：

def apply_model(X: tf.Tensor, my_model: tf.keras.Model) -> tf.Tensor:
        """
        Apply a tf.keras.Model to a tensor of unknown dimensions.

        Args:
            X (tf.Tensor): The tensor containing the input.
            my_model (tf.keras.Model): The model you want to apply.

        Returns:
            tf.Tensor: A tensor of the same shape as X, where all values are
                a prediction by the model.
        """
        if len(X.shape) > 1:
            result = tf.stack(
                [apply_model(x) for x in tf.unstack(X, axis=-1)],
                axis=-1,
            )
        else:
            result = my_model(X)

        return result

当然，您可以将其推广到模型采用多于一维的输入的情况。