Solved by ptrblck in post #14 
              
              
                Thanks for the code. 
Your return statement is at the wrong place. 
get_activation should return hook, not hook itself: 
Right way: 
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook

Your current implementation: 
def get_activat…