Save model when the hooks are used


I created a simple model, and registered the hooks to get the intermediate layer outputs,

    activation2 = OrderedDict()
    def get_activation2(name):
        def hook2(model, input, output):
            activation2[name] = output.detach()
        return hook2

    for name, layer in model.named_modules():

everthing is fine. When saving the model, using, 'model.pth'), I got this error -

AttributeError: Can’t pickle local object ‘get_activation2..hook2’

How to solve this, thanks

Same for me - any news?