Thanks for the feedback, it was None for sure and this caused the error.
Here is the full code of the problem.
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

n = 70 # num of points

# x is a tensor
x = torch.linspace(0, 10, steps=n)
k = torch.tensor(2.5)

# y is a tensor
y = k*x + 5*torch.rand(n)

# loss function
def mse(y_hat, y): return ((y_hat-y)**2).mean()

a = torch.tensor(-1., requires_grad=True)
a = nn.Parameter(a)

y_hat = a*x


plt.scatter(x, y);
plt.scatter(x,y_hat.detach());

lr = 0.1

def update_a():    
    global a
    loss = mse(y_hat, y) 
    print(loss) 
    loss.backward(a)
    print(a.grad)
    a.grad.zero_()
    a = a - lr * a.grad
   
    

for t in range(10): 
    a = update_a()


Plan is to learn the curve direction that was originally -1 and it should be 2.5.

image.png806×523 57.2 KB

Green should be more like blue dots.
I use the mse loss function, and I expect to learn the parameter a.