Hi everyone,
I’m tring to visualizing the feature map in my network. The code is as follow:
def hook_fn_backward(module, grad_input, grad_output):
    # print(module)
    # print('grad_output', grad_output)
    # print('grad_input', grad_input)
    total_grad_input.append(grad_input)
    total_grad_output.append(grad_output)

...
  PD2SE = PD2SEModel()

    # PD2SE = torch.load("./models/55.pth")
    PD2SE.to(device)
    modules = PD2SE.named_children()
    test = list(modules)
    for module in test[0][1]:
        if isinstance(module, nn.ReLU):
            module.register_backward_hook(hook_fn_backward)
...

By doing this, it exactally could output the feature map as I want(maybe not?) But I don’t know how to close it, it takes CUDA out of memory error. And I’ve read some documents about this, which says I should remove the hook. So my question is where should I remove it. Indeed, I think I could set a function save the feature map in local memory. Then I could close it since the code use the hook again. Looking forward your reply.
@albanD @ptrblck