A PyTorch hook that stores the state of a forward pass

stateful_hook.py

class StatefulHook:
    def __init__(self):
        self.module = None
        self.input = None
        self.output = None

    def __call__(self, module, input, output):
        self.module = module
        self.input = input
        self.output = output


net = ...  # some neural net
hook = StatefulHook()
net.my_layer.register_forward_hook(hook)

print(hook.input)
print(hook.output)