You have a small indentation error in your hook method:
activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook

The return statement should be called from get_activation.