Train a Linear Model

Write a real training loop in PyTorch.

Given X of shape (n, d) and y of shape (n,) as float tensors, build a torch.nn.Linear(d, 1), train it with plain SGD on mean squared error for epochs full batch steps, and return the trained module.

The loop is four lines and always in this order:

optimizer.zero_grad()
loss = criterion(model(X).squeeze(-1), y)
loss.backward()
optimizer.step()

zero_grad() first is the one people forget. Gradients accumulate in .grad by default, which is what makes gradient accumulation across mini-batches possible, and what silently ruins a loop that does not clear them: step two would descend the sum of two steps' gradients.

model(X) returns shape (n, 1) while y is (n,). Left alone, broadcasting turns the difference into an (n, n) matrix and the loss becomes meaningless while still producing a number. Squeeze the output.

This one is scored on the result, not on exact values: your trained model has to reach a mean squared error below 0.01 on the data it was given.

Function signature

def fit_linear(X, y, epochs, lr):

Example

Input

X = torch.tensor([[0.0], [1.0], [2.0]])
y = torch.tensor([1.0, 3.0, 5.0])        # y = 2x + 1
model = fit_linear(X, y, epochs=400, lr=0.05)
model(X).squeeze(-1)

Output

tensor([1.0000, 3.0000, 5.0000], grad_fn=<SqueezeBackward1>)

Reasoning

The data lies on y = 2x + 1. With zero_grad, backward and step in the right order, 400 SGD steps bring the weight to about 2 and the bias to about 1, so the predictions match y.

Constraints and notes

• Call optimizer.zero_grad() before each backward pass.
• Squeeze the model output from (n, 1) to (n,).
• The trained model must reach an MSE below 0.01.