Error in Adam Optimizer due to tf.assign function - tensorflow2.0

Question

I have reproduced the error from my code. As shown in the code, combination of "input" and "loss_total" gives the error where gradients are None type. How can I resolve this.
grads_and_vars = optimizer_utils.filter_empty_gradients(grads_and_vars)
File ~\AppData\Roaming\Python\Python39\site-packages\keras\optimizers\optimizer_v2\utils.py:77 in filter_empty_gradients
raise ValueError(
ValueError: No gradients provided for any variable: (['Variable_541:0', 'Variable_542:0'],). Provided grads_and_vars is ((None, <tf.Variable 'Variable_541:0' shape=(5, 5) dtype=float32>), (None, <tf.Variable 'Variable_542:0' shape=(5, 5) dtype=float32>)).
This function is used in the code.
def mean_squared_error(pred, exact):
if type(pred) is np.ndarray:
return np.mean(np.square(pred - exact))
else:
return tf.math.reduce_mean(tf.square(pred - exact))
class make_graph:
def __init__(self, ny, nx):
self.R = 287
self.u = tf.Variable(tf.zeros([ny+2, nx+2], tf.float32), trainable=True)
self.v = tf.Variable(tf.zeros([ny+2, nx+2], tf.float32), trainable=True)
self.U_padded = tf.Variable(tf.zeros([ny+2, nx+2], tf.float32), trainable=True)
self.V_padded = tf.Variable(tf.zeros([ny+2, nx+2], tf.float32), trainable=True)
#tf.function
def loss_estimation(self):
self.U_padded[1:-1,1:-1].assign(self.u)
self.V_padded[1:-1,1:-1].assign(self.v)
# self.loss_total = mean_squared_error(self.u-self.v+self.P-self.rho, 0)
self.loss_total = mean_squared_error(self.U_padded + self.V_padded, 0)
return self.loss_total
#tf.function
def train(self, optimizer):
with tf.GradientTape() as tape:
input = [self.u, self.v]
# input = [self.U_padded, self.V_padded]
self.loss = self.loss_estimation()
self.gradients = tape.gradient(self.loss, input)
optimizer.apply_gradients(zip(self.gradients, input))
if name == "main":
optimizer = tf.keras.optimizers.Adam(learning_rate = 0.001)
model = make_graph(3, 3)
for i in range(10):
model.train(optimizer)
print(i)