Thomas’ rule of thumb for when to use retain_grad: use it within one training step if you must, but not across steps.
If that doesn’t work for you, you’re either doing something beyond my wit or not quite doing it right yet.
Best regards
Thomas