Custom loss with weights when using BatchDataset in Keras

Question

I have a Keras model for which I have features, labels and an additional array which I want to use it as weights for a custom loss function. I am ingesting the data using 2 BatchDataset structures as follows:

1. One is containing the features and labels:

   train_ds.element_spec
   (TensorSpec(shape=(None, None, 100), dtype=tf.float32, name=None),
    TensorSpec(shape=(None, 100), dtype=tf.float32, name=None))
   

2. The additional BatchDataset with the weights for the loss function

   train_weights.element_spec
   TensorSpec(shape=(None, 100), dtype=tf.float32, name=None)
   

I am training the model this way:

model.fit(train_ds, epochs=10, batch_size=512, shuffle=True)

The custom loss function loss function which uses the additional data is like this:

def custom_loss(y_true, y_pred, weights):
    loss = (y_true - y_pred)*weights
    return loss

model.compile(loss=custom_loss, optimizer=optimizer, metrics=["mae"])

Is it possible to use the train_ds and train_weights BatchDatasets I have at the moment to define this custom loss function and use it for training?