Non-Linear regression model

I would like to create a model that predicts the equation of multiplying, division, or even power or roots the only equation that can be predicted is linear like ( pred = w * input + b )
so how to create that model with two input numbers that were multiplied

I tried to make a model that can predict the power of the input and I made a dataset with input x and output with y that’s y = x^2 .
I thought that if I separate the inputs 2 times for 2 neurons and get the log of each one I will get that log(x) and log(x).
then I collect the output at one neuron and get the exp(2log(x)) it will gives me x^2
so that I made this network.

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.fc1 = nn.Linear(1, 2)
        self.fc2 = nn.Linear(2, 2)
        self.fc3 = nn.Linear(2, 1)
    def forward(self, x):
        x = self.fc1(x)
        x = torch.log(x)
        x = self.fc2(x)
        x = torch.exp(x)
        x = self.fc3(x)
        return x

model = Network()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

and the predicted output is “tensor([[nan]])”
so how to make that network?

and another question I don’t know how to arrange the tensor for 2 neurons input of network.
how to arrange the shape of input tensors to make it acceptable by network of 2 or whatever number of inputs ?

It’s probably because you’re talking the log of the first layer, whose values could be negative, which leads to nans.