Here is the full workaround. (I think this query is more fit on S0).
from tensorflow.keras import Model, Input
import tensorflow as tf 
from tensorflow.keras import layers

latent_dim = 64

class Autoencoder(Model):
    def __init__(self, latent_dim):
        super(Autoencoder, self).__init__()
        self.latent_dim = latent_dim
        self.encoder = tf.keras.Sequential([
                                            layers.Flatten(),
                                            layers.Dense(latent_dim, activation='relu'),
                                            ])

        self.decoder = tf.keras.Sequential([
                                            layers.Dense(784, activation='sigmoid'),
                                            layers.Reshape((28, 28))
                                            ])

    def call(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

    def build_graph(self):
        x = Input(shape=(28,28))
        return Model(inputs=[x], outputs=self.call(x))

autoencoder = Autoencoder(latent_dim)

Model summary
autoencoder.build_graph().summary()
Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_5 (InputLayer)         [(None, 28, 28)]          0         
_________________________________________________________________
sequential (Sequential)      (None, 64)                50240     
_________________________________________________________________
sequential_1 (Sequential)    (None, 28, 28)            50960     
=================================================================
Total params: 101,200
Trainable params: 101,200
Non-trainable params: 0

Plotting
tf.keras.utils.plot_model(autoencoder.build_graph(), 
                          expand_nested=True, show_shapes=True)


download (1)512×777 16.2 KB