Function Call stack : train function - tensorflow

Question

I wrote a custom metric in tensorflow. It works fine during the first epoch. After , first epoch it generates an error. Where it says, "Function call stack:
train_function" Here is my code for , If i send pass to the reset state function, it works. So, far my guess is , the problem is in reset state function.
class PSNR(tf.keras.metrics.Metric):
def __init__(self, **kwargs):
super(PSNR, self).__init__(**kwargs)
self.psnr = self.add_weight(name='psnr', initializer='zeros')
self.num_ex = self.add_weight(name = 'num_ex', initializer='zeros')
self.total_psnr = self.add_weight(name = 'total_psnr', initializer='zeros')
def update_state(self, y_true, y_pred, sample_weight=None):
mse = tf.reduce_mean(tf.math.square(y_pred - y_true), axis = [1, 2, 3])
self.total_psnr.assign_add(tf.reduce_mean(4.34*tf.math.log(tf.math.divide(1, mse))))
self.num_ex.assign_add(1)
self.psnr = tf.math.divide(self.total_psnr, self.num_ex)
print(self.total_psnr)
print(self.psnr)
print(self.num_ex)
def result(self):
return self.psnr
def reset_state(self):
self.psnr = self.add_weight(name='psnr', initializer='zeros')
self.num_ex = self.add_weight(name = 'num_ex', initializer='zeros')
self.total_psnr = self.add_weight(name = 'total_psnr', initializer='zeros')
and the traceback is,
2022-08-30 04:46:57.091998: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2)
Epoch 1/2
2022-08-30 04:46:59.287913: I tensorflow/stream_executor/cuda/cuda_dnn.cc:369] Loaded cuDNN version 8005
44/44 [==============================] - 44s 797ms/step - loss: 0.0289 - psnr: 16.8205 - val_loss: 0.1307 - val_psnr: 9.2601
Epoch 2/2
2022-08-30 04:48:19.970412: W tensorflow/core/framework/op_kernel.cc:1692] OP_REQUIRES failed at resource_variable_ops.cc:548 : Not found: Resource localhost/_AnonymousVar54/N10tensorflow3VarE does not exist.
---------------------------------------------------------------------------
NotFoundError Traceback (most recent call last)
/tmp/ipykernel_17/4054115442.py in <module>
3 validation_data = valid_dataset,
4 #callbacks = [tfdocs.modeling.EpochDots()],
----> 5 verbose=1)
6 models.append(auto_en)
7 histories.append(history_auto_en)
/opt/conda/lib/python3.7/site-packages/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
1182 _r=1):
1183 callbacks.on_train_batch_begin(step)
-> 1184 tmp_logs = self.train_function(iterator)
1185 if data_handler.should_sync:
1186 context.async_wait()
/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
883
884 with OptionalXlaContext(self._jit_compile):
--> 885 result = self._call(*args, **kwds)
886
887 new_tracing_count = self.experimental_get_tracing_count()
/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
915 # In this case we have created variables on the first call, so we run the
916 # defunned version which is guaranteed to never create variables.
--> 917 return self._stateless_fn(*args, **kwds) # pylint: disable=not-callable
918 elif self._stateful_fn is not None:
919 # Release the lock early so that multiple threads can perform the call
/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in __call__(self, *args, **kwargs)
3038 filtered_flat_args) = self._maybe_define_function(args, kwargs)
3039 return graph_function._call_flat(
-> 3040 filtered_flat_args, captured_inputs=graph_function.captured_inputs) # pylint: disable=protected-access
3041
3042 #property
/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
1962 # No tape is watching; skip to running the function.
1963 return self._build_call_outputs(self._inference_function.call(
-> 1964 ctx, args, cancellation_manager=cancellation_manager))
1965 forward_backward = self._select_forward_and_backward_functions(
1966 args,
/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in call(self, ctx, args, cancellation_manager)
594 inputs=args,
595 attrs=attrs,
--> 596 ctx=ctx)
597 else:
598 outputs = execute.execute_with_cancellation(
/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
58 ctx.ensure_initialized()
59 tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60 inputs, attrs, num_outputs)
61 except core._NotOkStatusException as e:
62 if name is not None:
NotFoundError: Resource localhost/_AnonymousVar54/N10tensorflow3VarE does not exist.
[[node AssignAddVariableOp_3 (defined at tmp/ipykernel_17/1623473192.py:12) ]] [Op:__inference_train_function_1892]
Function call stack:
train_function