This would explain the error, since loss_iter1 would still reference the old computation graph.



 mbsyed:

Is there a way I can fix this?


It depends on your actual use case, since you are using stale forward activations to calculate the parameter updates, which is wrong.
From the linked issue:


loss_iter1 was calculated using the first forward pass and the model with parameter_set_0

this forward pass also calculated all intermediate forward activations (fwd_set_0) and stored them, which are needed to compute the gradients
you are updating the model to parameter_set_1, all forward activations are now stale, since they were not calculated by parameter_set_1


loss_iter1.backward() tries to compute the gradients using fwd_set_0 and parameter_set_1, which is wrong and fails

You could either delay the optimizer.step() call or recompute the forward activations depending on your use case.