Hi,
Thanks for the amazing solution. It works correctly for the forward pass.
I am using forward along with backward and doing certain operations. So, each time, when the backward is called, it creates a new thread and the output of backward is saved on that thread. My problem is how do I get the backward thread ID for the corresponding forward thread ID.
self.forward_hook = self.target_layer.register_forward_hook(self.hook_fn_act)
self.backward_hook = self.target_layer.register_backward_hook(self.hook_fn_grad)

def hook_fn_act(self, module, input, output):
    self.activations[threading.get_native_id()] = output.detach()

def hook_fn_grad(self, module, grad_input, grad_output):
    self.gradients[threading.get_native_id()] = grad_output[0].detach()

def foo():
    <Certain operations using the self.activations and self.gradients>
    <Problem is that thread id for the activations and gradients is not same, so how to know which 
      thread ID of activations maps to which one of gradients>