Loss & Training · Medium

PyTorch Training Step

Run one PyTorch forward, MSE loss, backward, optimizer step, and updated prediction pass.

Task

Implement linear_training_step(x: torch.Tensor, y: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, lr: float) -> torch.Tensor. Run exactly one PyTorch SGD training step for a bias-added linear regression model, then return the updated predictions.

Requirements

Treat weight and bias as initial values; do not mutate the caller's input tensors in-place.
Create trainable parameters from cloned weight and bias tensors.
Use predictions = x @ weight + bias.
Use mean squared error: torch.mean((predictions - y) ** 2).
Call optimizer.zero_grad(), loss.backward(), and optimizer.step() in the correct order.
Return predictions recomputed after the optimizer update.

Example

linear_training_step(
    torch.tensor([[1.0]]),
    torch.tensor([[2.0]]),
    torch.tensor([[0.0]]),
    torch.tensor([0.0]),
    0.1,
)
# tensor([[0.8000]])