Creating neural networks using tensorflow2

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt


class NN:
    def __init__(self):
        #  Create a hidden layer neural network model 
        self.model = tf.keras.Sequential()
        self.model.add(tf.keras.layers.Dense(10, activation='relu'))
        self.model.add(tf.keras.layers.Dense(1, activation='relu'))
        self.model.build((None, 2))     # 2 Express  input yes 2 dimension 

    def output(self, x):
        return self.model(x)

    def train(self, x, y):
        opt = tf.keras.optimizers.Adam()    # Adam Acceleration algorithm 
        loss = lambda: tf.reduce_mean(tf.square(y - self.output(x)))    #  Mean square error loss function 
        print(tf.reduce_mean(tf.square(y - self.output(x))).numpy())
        opt.minimize(loss, self.model.trainable_variables)


if __name__ == "__main__":
    nn = NN()
    #  Training set 
    x = np.float32(np.random.uniform(0, 1, size=(100, 1)))
    y = np.float32(np.random.uniform(0, 1, size=(100, 1)))
    z = 3 * x + y
    xx = np.hstack((x, y))
    for i in range(1000):
        nn.train(xx, z)

    #  Test set 
    x_verify = np.float32(np.random.uniform(0, 1, size=(100, 1)))
    y_verify = np.float32(np.random.uniform(0, 1, size=(100, 1)))
    z_verify = 3 * x_verify + y_verify
    xx_verify = np.hstack((x_verify, y_verify))
    z_ = nn.output(xx_verify).numpy()   # Tensor Tensor transformation numpy array

    plt.plot(z_verify)
    plt.plot(z_)
    plt.show()

tensorflow Official documents ：Module: tf  |  TensorFlow Core v2.9.1