Solved by rasbt in post #2 
              
              
                How about defining a hook? Would that solve your problem? E.g., something like 
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
 
        self.gradients = []

        ## your layers
        self.some_conv_layer.register_backward_hook(self.save_gradients)
   …