import copy
import torch
torch.manual_seed(42); torch.set_num_threads(1)
X=torch.randn(120,2)
y=2*X[:,0:1]-.5*X[:,1:2]
train_X,train_y=X[:80],y[:80]
val_X,val_y=X[80:100],y[80:100]
test_X,test_y=X[100:],y[100:]
model=torch.nn.Sequential(torch.nn.Linear(2,8),torch.nn.ReLU(),torch.nn.Linear(8,1))
optimizer=torch.optim.Adam(model.parameters(),lr=.03)
loss_fn=torch.nn.MSELoss()
initial=model[0].weight.detach().clone()
best=float('inf'); best_state=None
for epoch in range(160):
    model.train(); optimizer.zero_grad(set_to_none=True)
    pred=model(train_X)
    assert pred.shape==train_y.shape
    loss=loss_fn(pred,train_y)
    assert torch.isfinite(loss)
    loss.backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
    optimizer.step()
    model.eval()
    with torch.no_grad(): val=loss_fn(model(val_X),val_y).item()
    if val<best: best=val; best_state=copy.deepcopy(model.state_dict())
assert not torch.equal(initial,model[0].weight)
model.load_state_dict(best_state); model.eval()
with torch.no_grad(): test=loss_fn(model(test_X),test_y).item()
assert best_state is not None and test>=0
print('device=cpu','epochs=160','validation_mse=',round(best,6),'test_mse=',round(test,6))