Hey, the codebase was a bit of a mess to isolate a small sample, but I believe the following seems to have done the trick.
         # Multiply with logprobs
         generator_objective = (advantages * log_probs).sum(dim=0)
-        return (generator_objective, cumulative_rewards)
+        return (generator_objective, cumulative_rewards.clone())

I don’t have a requirement to call .backward() on the cumulative_rewards, and that seems to be creating an issue. Can you tell what’s happening under the hood?