A PyTorch hook that's registered in a `with` statement

context_hook.py

# This hook is particularly useful when ablating layers

class ContextHook:
    def __init__(self, layer):
        self.layer = layer

    def __enter__(self):
        self.handle = self.layer.register_forward_hook(self.hook)
        return self

    def __exit__(self, type, value, traceback):
        self.handle.remove()

    def hook(self, module, input, output):
        pass  # hook logic goes here


input = ...
net = ...  # some neural net
out = net(input)

with ContextHook(net.my_layer) as modified_net:
    modified_out = net(input)

print(out)
print(modified_out)