How to deal with incompatible shapes while trying to fit a model in Tensorflow Keras

Question

So, I have done some modifications with the VGG16 neural network to make a linear-regression and classification model. And at last I am writing the following code to compile and fit my data into the model to initiate training->

class Facetracker(Model):
#     The initialization function-
    def __init__(self,eyetracker,**kwargs):
        super().__init__(**kwargs)
        self.model = eyetracker #Instantiating the model

    def compile(self,opt,classlosss,localization_loss,**kwargs):
        super().compile(**kwargs)
        self.classloss = class_loss
        self.localization_loss = regress_loss
        self.opt = optimizer

        # Defining the training step
    def train_step(self,batch,**kwargs):
        X,y = batch #unpacking our data

        with tf.GradientTape() as tape:
            classes,coords = self.model(X,training=True)
            batch_classloss = self.classloss(y[0],classes)
            batch_localloss = self.localization_loss(tf.cast(y[1],tf.float32),coords)

            # calculating total loss-
            total_loss = batch_localloss+0.5*batch_classloss
            grad = tape.gradient(total_loss,self.model.trainable_variables)
        optimizer.apply_gradients(zip(grad,self.model.trainable_variables))

        return{
            "total_loss":total_loss,
            "class_loss":batch_classloss,
            "localilzation_loss":batch_localloss

        }
    def test_step(self,batch):
        X,y = batch
        classes,coords = self.model(X,training=False)
        batch_classloss = self.classloss(y[0],classes)
        batch_localloss = self.localization_loss(tf.cast(y[1],tf.float32),coords)

        total_loss = batch_localloss+0.5*batch_classloss

        return{
            "total_loss": total_loss,
            "class_loss": batch_classloss,
            "localilzation_loss": batch_localloss
        }
    # def call(self, X, **kwargs):
    #     return self.model(X,**kwargs)
# Replacing the call function with a lambda function
lambda self,X,**kwargs: self.model(X,**kwargs)

# Subclassing our model-
print("Subclassing.....")
model = Facetracker(facetracker)
print("Compiling......")
model.compile(optimizer,classlosss=class_loss,localization_loss=localization_loss)

# Preparing the log directory
logdir="logdir"
tensorboard_callbacks = tf.keras.callbacks.TensorBoard(log_dir=logdir)
print("Fitting the model")
hist = model.fit(train.take(80),
                 epochs=16,
                 initial_epoch =8,
                 validation_data=val,
                 validation_steps =8,
                 validation_freq=2,
                 callbacks = [[tensorboard_callbacks]])

This includes the class prepared for training the model and the last few lines are subclassing the model and fitting the prepared data into the model. The error I now get seems pretty hefty to me and it goes like this->

  File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\main.py", line 535, in <module>
    hist = model.fit(train.take(80),
  File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\MarkATT\lib\site-packages\keras\utils\traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\MarkATT\lib\site-packages\tensorflow\python\eager\execute.py", line 54, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:

Detected at node 'gradient_tape/sub_2/BroadcastGradientArgs' defined at (most recent call last):
    File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\main.py", line 535, in <module>
      hist = model.fit(train.take(80),
    File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\MarkATT\lib\site-packages\keras\utils\traceback_utils.py", line 64, in error_handler
      return fn(*args, **kwargs)
    File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\MarkATT\lib\site-packages\keras\engine\training.py", line 1409, in fit
      tmp_logs = self.train_function(iterator)
    File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\MarkATT\lib\site-packages\keras\engine\training.py", line 1051, in train_function
      return step_function(self, iterator)
    File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\MarkATT\lib\site-packages\keras\engine\training.py", line 1040, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\MarkATT\lib\site-packages\keras\engine\training.py", line 1030, in run_step
      outputs = model.train_step(data)
    File "C:\Users\Radhe Krishna\OneDrive\Documents\MarkATT\main.py", line 498, in train_step
      grad = tape.gradient(total_loss,self.model.trainable_variables)
Node: 'gradient_tape/sub_2/BroadcastGradientArgs'
Incompatible shapes: [2,4] vs. [8]
         [[{{node gradient_tape/sub_2/BroadcastGradientArgs}}]] [Op:__inference_train_function_22026]

It would be a great help if someone could examine this problem and help me get out of it. I have been quite struggling with it now. Thanks in advance!!!

I has planned to initiate training with the code provided, but I firstly received an error for the call function and as I resolved it by converting it into lambda function- This issue of incompatible shapes came up. I tries to adjust the integers in the input data and batch related fields but nothing worked.