tensorflow Exception encountered when calling layer (type CategoryEncoding)

I'm trying to code a layer to interface between a data set (numerical and categorical features) so it can be fed into a model. I can't understand the error I get when it comes to categorical columns.

ValueError: Exception encountered when calling layer (type CategoryEncoding). When output_mode is not 'int', maximum supported output rank is 2. Received output_mode multi_hot and input shape (10, 7, 1), which would result in output rank 3.

From what I understand, the batch size should not have been counted in, but it is. And that seems to break. Note that reproducing with only numerical features works fine.

Thank you for your help.

import tensorflow as tf
import pandas as pd
import numpy as np

# Simulate a data set of categorical and numerical values
# Configure simulation specifications: {feature: number of unique categories or None for numerical}
theSimSpecs = {'Cat1': 54, 'Cat2': 2, 'Cat3': 4, 'Num1': None, 'Num2': None}
# theSimSpecs = {'Num1': None, 'Num2': None}

# batch size and timesteps
theBatchSz, theTimeSteps = 10, 4

# Creation of the dataset as pandas.DataFrame
theDFs = []
for theFeature, theUniques in theSimSpecs.items():
    if theUniques is None:
        theDF = pd.DataFrame(np.random.random(size=theBatchSz * theTimeSteps), columns=[theFeature])
    else:
        theDF = pd.DataFrame(np.random.randint(low=0, high=theUniques, size=theBatchSz * theTimeSteps),
                             columns=[theFeature]).astype('category')
    theDFs.append(theDF)
theDF = pd.concat(theDFs, axis=1)

# code excerpt
# inventory of the categorical features' values ( None for the numerical)
theCatCodes = {theCol: (theDF[theCol].unique().tolist() if str(theDF[theCol].dtypes) == "category" else None)
               for theCol in theDF.columns}

# Creation of the batched tensorflow.data.Dataset
theDS = tf.data.Dataset.from_tensor_slices(dict(theDF))
theDS = theDS.window(size=theTimeSteps, shift=1, stride=1, drop_remainder=True)
theDS = theDS.flat_map(lambda x: tf.data.Dataset.zip(x))
theDS = theDS.batch(batch_size=theTimeSteps, drop_remainder=True)
theDS = theDS.batch(batch_size=theBatchSz, drop_remainder=True)

# extracting one batch
theBatch = next(iter(theDS))
tf.print(theBatch)

# Creation of the components for the interface layer
theFeaturesInputs = {}
theFeaturesEncoded = {}

for theFeature, theCodes in theCatCodes.items():
    if theCodes is None: # Pass-through for numerical features
        theNumInput = tf.keras.layers.Input(shape=[], dtype=tf.float32, name=theFeature)
        theFeaturesInputs[theFeature] = theNumInput

        theFeatureExp = tf.expand_dims(input=theNumInput, axis=-1)
        theFeaturesEncoded[theFeature] = theFeatureExp

    else: # Process for categorical features
        theCatInput = tf.keras.layers.Input(shape=[], dtype=tf.int64, name=theFeature)
        theFeaturesInputs[theFeature] = theCatInput

        theFeatureExp = tf.expand_dims(input=theCatInput, axis=-1)
        theEncodingLayer = tf.keras.layers.CategoryEncoding(num_tokens=theSimSpecs[theFeature], name=f"{theFeature}_enc",
                                                            output_mode="multi_hot", sparse=False)
        theFeaturesEncoded[theFeature] = theEncodingLayer(theFeatureExp)

theStackedInputs = tf.concat(tf.nest.flatten(theFeaturesEncoded), axis=1)

theModel = tf.keras.Model(inputs=theFeaturesInputs, outputs=theStackedInputs)

theOutput = theModel(theBatch)
tf.print(theOutput)