tensorflow: tf-nightly-cpu couldn't trace any graph with subclass models.

System information

• Have I written custom code: Yes.
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Win 10.
• TensorFlow installed from (source or binary): pip intall tf-nightly-cpu.
• TensorFlow version (use command below): 2.2.0-dev20200426.
• Python version: 3.7.7
• Tensorboard version: 2.3.0a20200412 or 2.1.1

Describe the current behavior

I’ve created a subclass model, and trained it with tensorboard callback. Then I start tensorboard and select Graphs dashboard, but it shows error messages that Graph visualization failed as the picture shows：

I’ve checked the issue 1961 here but didn’t get any help.

By the way, Sequential Model could trace the graph.

Describe the expected behavior

Things works well when I use tensorflow 2.1.0 with both tensorboard 2.3.0a20200412 and 2.1.1. So I think tf-nightly-cpu should work as well.

Standalone code to reproduce the issue

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np


class ThreeLayerMLP(keras.Model):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.dense_1 = layers.Dense(64, activation='relu', name='dense_1')
        self.dense_2 = layers.Dense(64, activation='relu', name='dense_2')
        self.pred_layer = layers.Dense(10, name='predictions')

    def call(self, inputs):
        x = self.dense_1(inputs)
        x = self.dense_2(x)
        return self.pred_layer(x)


model = ThreeLayerMLP(name='3_layer_mlp')

x_train, y_train = (np.random.random(
    (60000, 784)), np.random.randint(10, size=(60000, 1)))
x_test, y_test = (np.random.random(
    (10000, 784)), np.random.randint(10, size=(10000, 1)))

model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.RMSprop())

callback = tf.keras.callbacks.TensorBoard(
    'subclass_logs',
    update_freq=2,
)
history = model.fit(x_train,
                    y_train,
                    batch_size=64,
                    epochs=10,
                    callbacks=[callback])