Hi Zeeky
Since model parameters in PyTorch are all leaf variables, which means you can not just set the parameters to new values manually. To achieve this, you just need to wrap the update function in a torch.no_grad() block:
with torch.no_grad():
    param = param - lr * grad