Hi @ptrblck, I have a similar but different problem. I want to calculate the loss2 based on params s1, and get the gradients w.r.t s0, then update s0 with the gradients. But the there is a problem when calculating the gradients, which was expressed in the link. Could you please give me some suggestions? Thank you.