I hope that means my solution is correct! 
I was wondering if I could ask a follow up question? The register_full_backward_hook if I call a double backward on my network does the backward_hook recorded a double backward or is it just a single backward? And, if I call 2 in a row (i.e. a double backward with one loss function, then a single backward with a different loss function). Which output is stored?
my forward/backward hooks are defined like so,
  def _save_input(self, module, input):
    self.state[module]['a'] = input[0]
    
  def _save_output(self, module, grad_input, grad_output):
    self.state[module]['gl'] = grad_output[0]

So, I’d assume whenever a backward call is called, it overwrites the previous call?
Thank you for all the help!