pytorch_hook.py

class SaveFeatures():
    def __init__(self, module):
        self.hook = module.register_forward_hook(self.hook_fn)
    def hook_fn(self, module, input, output):
        self.features = torch.tensor(output,requires_grad=True).cuda()
    def close(self):
        self.hook.remove()

Thanks for your article about neural net visualization! For more recent versions of PyTorch, I found it is necessary to change line 5 to self.features = output for the back-propagation to succeed.