Pytorch Grad!

import torch

# We want to find a 'w' such that y is close to w * x.
# The true 'w' is clearly 2.
X = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32)
Y = torch.tensor([2.0, 4.0, 6.0, 8.0], dtype=torch.float32)

# Initialize our weight 'w' with a random guess.
# requires_grad=True tells PyTorch to track this tensor for gradient calculations.
w = torch.tensor(0.0, dtype=torch.float32, requires_grad=True)

def forward(x):
    return w * x

def loss(y, y_predicted):
    return ((y_predicted - y)**2).mean()

learning_rate = 0.01
n_iters = 20

print(f"Starting training... Initial weight w = {w.item():.3f}")

for epoch in range(n_iters):
    # Forward pass
    y_pred = forward(X)

    # loss
    l = loss(Y, y_pred)

    # Calculate gradients = backward pass
    # core of autograd. It calculates the derivative of 'l'
    # with respect to every tensor that has requires_grad=True (i.e., 'w').
    l.backward() # dl/dw

    # Manually update weights
    # torch.no_grad() because this is not part of the computation graph
    with torch.no_grad():
        # The calculated gradient is now in w.grad
        # Update rule: w = w - learning_rate * gradient
        w.data -= learning_rate * w.grad

    # Zero out the gradients for the next iteration
    if (epoch + 1) % 2 == 0:
        print(f'Epoch {epoch+1}: w = {w.item():.3f}, w gradient: {w.grad}, loss = {l.item():.8f}')

    w.grad.zero_()

    # if (epoch + 1) % 2 == 0:
    #     print(f'Epoch {epoch+1}: w = {w.item():.3f}, w gradient: {w.grad}, loss = {l.item():.8f}')

print(f"\nTraining finished. The learned weight is: {w.item():.3f}")
print("The true weight was 2.0")

Starting training... Initial weight w = 0.000
Epoch 2: w = 0.555, w gradient: -25.5, loss = 21.67499924
Epoch 4: w = 0.956, w gradient: -18.423751831054688, loss = 11.31448650
Epoch 6: w = 1.246, w gradient: -13.311159133911133, loss = 5.90623236
Epoch 8: w = 1.455, w gradient: -9.61731243133545, loss = 3.08308983
Epoch 10: w = 1.606, w gradient: -6.948507308959961, loss = 1.60939169
Epoch 12: w = 1.716, w gradient: -5.020296096801758, loss = 0.84011245
Epoch 14: w = 1.794, w gradient: -3.627163887023926, loss = 0.43854395
Epoch 16: w = 1.851, w gradient: -2.6206254959106445, loss = 0.22892261
Epoch 18: w = 1.893, w gradient: -1.8934016227722168, loss = 0.11949898
Epoch 20: w = 1.922, w gradient: -1.3679819107055664, loss = 0.06237914

Training finished. The learned weight is: 1.922
The true weight was 2.0