Hi,
If you do new_param = old_param + lr * grad then it will be properly differentiable.
Note that if you try ot work with torch.nn though, the nn.Parameter can only be leafs (and so cannot have history). So you won’t be able to use them here.