Get batch gradient for each individual element using hook

Hi,

I am trying to implement Grad-CAM with batch and I am referring to thisGrad-CAM repository. To my understanding, I just need to modify the shape of the one_hot_output (the first dimension) and accordingly get the gradients for each individual element.

Currently, the issue is that although I can get the right shape of each variable such as gradients and weights the final activation map seems incorrect when I was overlaying the CAM into the original images.