The loss computation looks a bit concerning, as it seems you might be recreating tensors and thus detaching them from the computation graph.
What is the type of batch_logp and batch_weights? Is one of these tensors the output of both models (called sequentially)?
Also, could you post the code of compute_loss? If you are detaching the tensors, the batch_loss.backward() call should raise an error, so I’m not sure where exactly the issue comes from.