Malformed Op graph of custom Keras model

While running a custom Keras model with tensorboard callback. The Conceptual graph is generated, however, the Op graph returns: Error: Malformed GraphDef. I tried some existing suggestions related to potential naming conflicts and using name_scope, however, to no avail.

logdir = 'logs/func/' +"%Y%m%d-%H%M%S")

class MCLayer(tf.keras.layers.Layer):
  def __init__(self, name=None):
    super(MCLayer, self).__init__(name=name)
    #with tf.name_scope('test1'): 
    self.nT = tf.constant(400)
    self.n = tf.constant(100000)
    self.dt = tf.constant(1/365)
    self.drift = tf.constant(0.08)
    self.sigma = tf.constant(0.1)

  def call(self, inputs):
    #with tf.name_scope('test2'):    
    dWt = tf.random.normal(mean=0, stddev=tf.math.sqrt(self.dt), shape=[self.nT, self.n])
    dYt = self.drift*self.dt + self.sigma*dWt
    C = tf.cumsum(dYt, axis=0)
    S = tf.exp(C)
    A = tf.reduce_mean(S, axis=0)
    P = tf.reduce_mean(tf.maximum(A - inputs, 0))
    return P

input_layer = tf.keras.layers.Input(shape=(1), name='input_layer')
output_layer = MCLayer(name='output_layer')(input_layer)

model = tf.keras.models.Model(input_layer, output_layer, name='SomeModel')
result = model.predict(tf.constant(1.0, shape=(1,)), callbacks=[tf.keras.callbacks.TensorBoard(log_dir=logdir)])

I currently faced the same problem with TF2.8 with model.predict
Is there any problem with the TF version?

Refer to the issue: Graph visualization failed: GraphDefs failed to reconcile. · Issue #1929 · tensorflow/tensorboard · GitHub

I’ve tried some of the suggestions from issue #1929 such as using tf.name_scope('test1') . You can see those parts as comments in my code as it didn’t solve the problem.

Do you know what else might help?