Getting Incompatible shapes error with batch_size>1 with keras

I recently started working with keras and have looked at every available solution but its did not work. I have a simple neural network to compute a tensor value from 5 input constants, and 10 input tensors.This is my neural network, the main input layer inputs 5 constants and tensor input layer inputs 10 tensors:

import numpy as np
from keras.utils import plot_model
from keras.models import Model
from keras.layers import Input,Lambda
from keras.layers import Dense
from keras import backend as K


def function(x): #This function is used for the last layer to compute Anisotropic R-S
    tensor = x[0]
    constants = x[1]
    a = K.zeros(shape=(3,3,))

    for i in range(10):
        a = a + constants[:,i]*tensor[:,:,:,i]
    return a


main_input = Input(shape = (5,),name = 'main_input') #The invariant inputs
hidden1 = Dense(10,activation = 'relu')(main_input)
hidden2 = Dense(10,activation= 'relu')(hidden1) #10 constants

tensor_input = Input(shape= (3,3,10,),name = 'tensor_input')

output_layer = Lambda(function)([tensor_input,hidden2])

model = Model(inputs = [main_input,tensor_input], outputs = output_layer)
print(model.summary())

model.compile(optimizer = 'adam', loss='mean_squared_error', metrics=['accuracy'])
plot_model(model, to_file='multilayer_perceptron_graph.png')

#Just test inputs and outputs to correct shape
I_test = np.ones((120,5))
T_test = np.ones((120,3,3,10))
a_test = np.ones((120,3,3))

model.fit({'main_input': I_test, 'tensor_input': T_test},a_test,epochs=50,batch_size=2)

I have used a lambda layer as the output layer. It computes the tensor a as: a = g1*T1 + g2*T2 + ....g10*T10 where g's are constants which are computed from Dense 24 layer. There is no activation here, its simple linear combination. The output a is (3,3) matrix. The input tensors are 10 3*3 tensors, thus shape of I is (129,3,3,10).

I am getting the following error whenever batchsize>1:

InvalidArgumentError: Incompatible shapes: [2] vs. [2,3,3]
     [[Node: training_7/Adam/gradients/lambda_12/mul_9_grad/BroadcastGradientArgs = BroadcastGradientArgs[T=DT_INT32, _class=["loc:@training_7/Adam/gradients/lambda_12/mul_9_grad/Reshape"], _device="/job:localhost/replica:0/task:0/device:CPU:0"](training_7/Adam/gradients/lambda_12/mul_9_grad/Shape, training_7/Adam/gradients/lambda_12/mul_9_grad/Shape_1)]]

Please help me solve this problem.