Computing Jacobian and Derivative in Tensorflow is extremely slow

Is there a more efficient way to compute Jacobian (there must be, it doesn't even run for a single batch) I want to compute the loss as given in the self-explanatory neural network. Input has a shape of (32, 365, 3) where 32 is the batch size. The loss I want to minimize is Equation 3 of the paper.

I believe that I am not using the GradientTape optimally.

def compute_loss_theta(tape, parameter, concept, output, x):

    b = x.shape[0]
    in_dim = (x.shape[1], x.shape[2])

    feature_dim = in_dim[0]*in_dim[1]

    J = tape.batch_jacobian(concept, x)
    
    grad_fx = tape.gradient(output, x)
    grad_fx = tf.reshape(grad_fx,shape=(b, feature_dim))
    J = tf.reshape(J, shape=(b, feature_dim, feature_dim))

    parameter = tf.expand_dims(parameter, axis =1)

    loss_theta_matrix = grad_fx - tf.matmul(parameter, J)

    loss_theta = tf.norm(loss_theta_matrix)

    return loss_theta


for i in range(10):
    for x, y in train_dataset:

        with tf.GradientTape(persistent=True) as tape:
            tape.watch(x)
            
            parameter, concept, output = model(x)

            loss_theta = compute_loss_theta(tape, parameter, concept, output , x)

            loss_y = loss_object(y_true=y, y_pred=output)
                
            loss_value = loss_y + eps*loss_theta

        gradients = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))