Regression for a naive example

Hi,
I got a problem when I tried a naive regression example on 4 points. Here is my code.

import numpy as np
import torch
from torch import nn
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

model = nn.Sequential(
    nn.Linear(D, H),
    nn.Sigmoid(),
    nn.Linear(H, C)
)
model.to(device) # Convert to CUDA
criterion = torch.nn.MSELoss()

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) # built-in L2
# Training
for t in range(10000):
    # Feed forward to get the logits
    y_pred = model(x)
    # Compute the loss (MSE)
    loss = criterion(y_pred, y)
    print("[EPOCH]: %i, [LOSS or MSE]: %.6f" % (t, loss.item()))
    display.clear_output(wait=True)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
D = 1 # or 2
H = 30 # for example
C = 1
#--------------------
x1 = torch.tensor(np.array([np.linspace(-1,1,4)]).T, dtype=torch.float).to(device)
y = torch.cos(1000*x1)

x2 = torch.tensor(np.array([np.linspace(-1,1,4), np.linspace(-1,1,4)]).T, dtype=torch.float).to(device)
y = torch.cos(1000*x2[:,0])

When I give X1 (in this case D=1), the loss converges to 0. But when I give X2 (which is just X1 repeated twice and in this case D=2), the solver get stuck and cannot be fitted correctly. I tried with SGD and some other learning_rates also but did not work. What is the issue ?

Thanks

In the second use case you are ignoring the warning:

UserWarning: Using a target size (torch.Size([4])) that is different to the input size (torch.Size([4, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.

which explains that broadcasting will be used. In your case the wrong loss calculation will thus be used and your training gets stuck. Unsqueeze dim1 in y and it should work:

x2 = torch.tensor(np.array([np.linspace(-1,1,4), np.linspace(-1,1,4)]).T, dtype=torch.float).to(device)
y = torch.cos(1000*x2[:,0])
y = y.unsqueeze(1)

It works now ! Thanks a lot !