TypeError: call() got an unexpected keyword argument 'use_causal_mask' ---> getting this error on flickr8k/flickr30k dataset

Error

TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_23/1382744270.py in <module>
2 image_path = tf.keras.utils.get_file('surf.jpg', origin=image_url)
3 image = load_image(image_path)
----> 4 model.simple_gen(image, temperature=1, max_run=1)
/tmp/ipykernel_23/644194926.py in simple_gen(self, image, temperature, max_run)
7   for n in range(max_run):
8     print("interation", n)
----> 9     preds = self((img_features, tokens)).numpy()  # (batch, sequence, vocab)
10     preds = preds[:,-1, :]  #(batch, vocab)
11     if temperature==0:
/opt/conda/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
1035         with autocast_variable.enable_auto_cast_variables(
1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
1038 
1039         if self._activity_regularizer:
/tmp/ipykernel_23/4251248884.py in call(self, inputs)
20     # Look at the image
21     for dec_layer in self.decoder_layers:
---> 22         txt = dec_layer(inputs=(image, txt))
23 
24     txt = self.output_layer(txt)
/opt/conda/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
1035         with autocast_variable.enable_auto_cast_variables(
1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
1038 
1039         if self._activity_regularizer:
/tmp/ipykernel_23/725126540.py in call(self, inputs, training)
16 
17     # Text input
---> 18     out_seq = self.self_attention(out_seq)
19 
20     out_seq = self.cross_attention(out_seq, in_seq)
/opt/conda/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
1035         with autocast_variable.enable_auto_cast_variables(
1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
1038 
1039         if self._activity_regularizer:
/tmp/ipykernel_23/1379568068.py in call(self, x)
8 
9   def call(self, x):
---> 10     attn = self.mha(query=x, value=x, use_causal_mask=True)
11     print("mask", attn, "end_mask")
12     x = self.add([x, attn])
/opt/conda/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
1035         with autocast_variable.enable_auto_cast_variables(
1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
1038 
1039         if self._activity_regularizer:
TypeError: call() got an unexpected keyword argument 'use_causal_mask'

Code:

tensorflow tutorial for image captioning (instead of mobilenetv3, I am using inceptionV3)

dataset:

filckr8k dataset from kaggle -> https://www.kaggle.com/datasets/adityajn105/flickr8k

The error is generated from multiheadattention of causal self-attention layer
I have tried removing the use_causal_mask kwarg from call method and it seems to work fine

use_causal_mask was introduced in Tensorflow version 2.10.0.
Read more in this blog

Verify that you are using tensorflow 2.10.0 or above.