Your code has a small issue with the shapes.
If you are using
nn.MSELoss the shapes of your output and target should be the same.
However, in your current code the target is missing dim1, which will trigger some unwanted broadcasting inside the loss function.
train_x = torch.stack((torch.linspace(0, 10, 1000), torch.linspace(0.5, 10.5, 1000), torch.linspace(3, 14, 1000), torch.linspace(7.5, 19.5, 1000)),1)
train_y = (train_x[:,3]**2 - train_x[:,2])/10 + torch.cos(train_x[:,0] * (2 * math.pi)) + torch.sin(train_x[:,1] * (2 * math.pi)) + 0.1 * torch.randn(1000)
After fixing this issue, I played around with your code and your model is able to learn an approximation of the curve using:
self.base = nn.Sequential(
self.base_linear = nn.Linear(2, 1)
def forward(self, inputs):
x = inputs
hidden_base = self.base(x)
The model learned the approximation also using your current architecture, but was quite sensible to the random seed (a lot of solutions just learned the constant horizontal line). This modification is also not very robust, but might be a good starter to further experiment with your code.