Solved by albanD in post #8 
              
              
                If you use nightly build, you can pass inputs= to .backward() to specify what you want to compute gradients for. In this case ds_loss.backward(inputs=(lr_list[j],), retrain_graph=True). 
Otherwise, you can do grad, = autograd.grad(ds_loss, lr_list[j], retain_graph=True); lr_list[j].grad = grad but …