Save model when the hooks are used

Hello,

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():
        layer.register_forward_hook(get_activation2(name))

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

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

How to solve this, thanks

1 Like

Same for me - any news?