I am totally new to Pytorch and machine learning. I am trying to construct my model from scratch. The model is not CNN or RNN, just based on my formula. The input is two matrixes. What I want to do in my hidden layer is multiplying these two matrixes, and then output the result in the output layer.
class myModel(nn.Module):
    def __init__(self, matrix_a, matrix_b):
        super(myModel, self).__init__()
        self.matrix_a = matrix_a
        self.matrix_b = matrix_b
        self.build_model()


    def build_model(self):
        self.layers = nn.ModuleList()
        #i2h
        i2h = self.build_input_layer()
        self.layers.append(i2h)
        # h2h
        h2h = self.build_hidden_layer()
        self.layers.append(h2h)
        # h2o
        h2o = self.build_output_layer()
        self.layers.append(h2o)


    def build_input_layer(self):
        input_matrix_a = self.matrix_a
        input_matrix_b = self.matrix_b
        return


    def build_hidden_layer(self):
        # I want to multiply matrix a and b
        return


    def forward(self,matrix_a, matrix_b):
        for layer in self.layers:
            output = layer(matrix_a, matrix_b) # not sure if it is correct
        return output