I am trying to run the keras-tutorial Probabilistic Bayesian Neural Networks to get an understanding of Bayesian neural networks (BNN). The tutorial contains a google-colab notebook, so you can run it directly within the browser. However, when I try to train the BNN for the first time (chapter Train BNN with a small training subset), it throws this error:

AttributeError                            Traceback (most recent call last)
 in ()
      3 small_train_dataset = train_dataset.unbatch().take(train_sample_size).batch(batch_size)
      4 
----> 5 bnn_model_small = create_bnn_model(train_sample_size)
      6 run_experiment(bnn_model_small, mse_loss, small_train_dataset, test_dataset)

2 frames
/usr/local/lib/python3.10/dist-packages/tf_keras/src/engine/input_spec.py in assert_input_compatibility(input_spec, inputs, layer_name)
    249                 )
    250         if spec.min_ndim is not None:
--> 251             ndim = x.shape.rank
    252             if ndim is not None and ndim < spec.min_ndim:
    253                 raise ValueError(

AttributeError: 'tuple' object has no attribute 'rank'

The error message probably refers to the model definition in the cell before:

Model Definition

def create_bnn_model(train_size):
    inputs = create_model_inputs()
    features = keras.layers.concatenate(list(inputs.values()))
    features = layers.BatchNormalization()(features)

    # Create hidden layers with weight uncertainty using the DenseVariational layer.
    for units in hidden_units:
        features = tfp.layers.DenseVariational(
            units=units,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1 / train_size,
            activation="sigmoid",
        )(features)

    # The output is deterministic: a single point estimate.
    outputs = layers.Dense(units=1)(features)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model

The code to start the training is this:
Start training

num_epochs = 500
train_sample_size = int(train_size * 0.3)
small_train_dataset = train_dataset.unbatch().take(train_sample_size).batch(batch_size)

bnn_model_small = create_bnn_model(train_sample_size)
run_experiment(bnn_model_small, mse_loss, small_train_dataset, test_dataset)

I tried wrapping the features - tensor into a tf.keras.input - layer, but it did not solve the problem