Thanks Peter for your reply.
Using retain_graph=True on the full graph seemed unnecessary, but your reply from a different thread helped me resolve the issue: https://discuss.pytorch.org/t/how-to-free-graph-manually-after-using-retain-graph-true/66258/2