@Rahul_Chand Thank you very much for the detailed analysis and for explaining how you were able to debug the issue! I understand now where the problem comes from, the zero_ inplace operation does not make autograd “forget” the previous history even though it will no longer be relevant for computing the gradient value.
In your simplified example, I could fix this problem by replacing the in-place operation with the following: self.x[0] = o1
However, I have not been able to to transfer this to my version where I have a wrapped model with multiple Parameters inside where I don’t know the exact structure (wanting to keep this generic), so can only access it through either actual_model.parameters() or actual_model.state_dict().
Do you have any idea how I could achieve this?