examples/demo_subclass.py/MyModel/__init__
class MyModel:    def __init__(self, hidden_dim, output_dim):
        super().__init__()
        self.dense1 = layers.Dense(hidden_dim, activation="relu")
        self.dense2 = layers.Dense(hidden_dim, activation="relu")
        self.dense3 = layers.Dense(output_dim)