I’m writing an implementation of the GradCAM algorithm to visualize my neural net. In order to achieve this I’m making use of hooks as so:
...
def save_gradients(self, in):
    self.gradients = in

...

def forward(self, x)
    if i == target_layer:
        x.register_hook(self.save_gradients)
        out = x
...

The problem is this only sometimes works, depending on what layer I’m trying to look at. For example:
Sequential(                                                                                                                                          
  (Conv2d): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)                                                                                                                    
  (BatchNorm2d): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)                                                                                              
  (activation): LeakyReLU(negative_slope=0.1, inplace=True)
)

will work, but not:
Sequential(                                                                                                                                [124/1941]
  (Conv2d): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)                                                                                                   
  (BatchNorm2d): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)                                                                                             
  (activation): LeakyReLU(negative_slope=0.1, inplace=True)                                                                                                                                    
)

Is my intuition/knowledge wrong about something, hence why I can only look at certain layers, or is my code wrong elsewhere?