Keras multi-label classification with binary cross_entropy loss: invalid shape

Question

I am working on a multi-label classification task with a tf dataset that looks like the following:

<PrefetchDataset shapes: ({cluster_6: (None, 1), cluster_3: (None, 1), cluster_10: (None, 1), cluster_4: (None, 1), request_time: (None, None), userId: (None, None), cluster_5: (None, 1), cluster_7: (None, 1), cluster_8: (None, 1), cluster_9: (None, 1), cluster_1: (None, 1), cluster_2: (None, 1)}, (None, 5)), types: ({cluster_6: tf.float32, cluster_3: tf.float32, cluster_10: tf.float32, cluster_4: tf.float32, request_time: tf.int64, userId: tf.int64, cluster_5: tf.float32, cluster_7: tf.float32, cluster_8: tf.float32, cluster_9: tf.float32, cluster_1: tf.float32, cluster_2: tf.float32}, tf.float32)>

Here cluster_1, ... cluster_10 are my features, and my target is a binary vector of dimension 5 as highlighted above. (i.e. there are 5 classes)

I select the features by a custom layer: CustomDenseFeatures(tf.keras.layers.Layer)

and construct a deep and cross network as follows

def get_dcn_outputs():
    inputs = KERAS_INPUTS.copy()

    dense_feats_layer = CustomDenseFeatures()
    dense_feat_names = dense_feats_layer.feature_names
    dense_feats = dense_feats_layer({
        k: v for k, v in inputs.items() if k in dense_feat_names
    })
            
    dense_feats = dense_feats['cont']
    
    dcn_stacked_output = DeepAndCrossNetwork(
        poly_degree=DCN_POLY_DEGREE,
        dnn_layer_sizes=DCN_LAYER_SIZES,
    )(dense_feats)
    
    dcn_logit = tf.keras.layers.Dense(units=len(LABEL_NAME), name='dcn_logit')(dcn_stacked_output) # unit = dimension of output

    return tf.keras.layers.Activation(
        tf.keras.activations.sigmoid,
        name='label'
    )(dcn_logit))

I then compile the model with

def get_compiled_model():
    outputs = get_dcn_outputs()
    model = tf.keras.Model(inputs=KERAS_INPUTS, outputs=outputs)
    model.compile(
        optimizer=tfa.optimizers.LazyAdam(learning_rate=INITIAL_LEARNING_RATE, clipnorm=CLIPNORM),
        loss=tf.keras.losses.BinaryCrossentropy(from_logits=False),
        weighted_metrics=[
            Calibration(),
            RelativeCoefficientOfDiscrimination(),
            RelativeCrossEntropy(baseline_pred=LABEL_PRIOR.numpy()),
            SigmoidScaledAUC(
                num_thresholds=500, name='roc_auc', curve='ROC', baseline_pred=LABEL_PRIOR.numpy()
            ),
            SigmoidScaledAUC(
                num_thresholds=500, name='pr_auc', curve='PR', baseline_pred=LABEL_PRIOR.numpy()
            ),
        ]
    )
    return model

and the model summary shows

Now I try to fit the model. Unfortunately I am seeing

ValueError: Can not squeeze dim[1], expected a dimension of 1, got 5 for '{{node binary_crossentropy_1/weighted_loss/Squeeze}} = Squeeze[T=DT_FLOAT, squeeze_dims=[-1]](ones_like_2)' with input shapes: [?,5].

I think the error is related to the weighted_metrics, which is a of metrics to be evaluated and weighted by sample_weight or class_weight during training and testing. But I am not sure. Appreciate if anyone can help with it.