castorini/daam

Trace hook never unhooks the attention processor

Closed this issue · 0 comments

daam/daam/trace.py

Lines 281 to 282 in a323129

def _hook_impl(self):
self.module.set_processor(self)

This hook doesn't unhook which can cause problems for the network in the future.

So I added a capture of the original processor and then added an unhook impl to set it back to the original processor.

    def _hook_impl(self):
        self.original_processor = self.module.processor
        self.module.set_processor(self)

    def _unhook_impl(self):
        self.module.set_processor(self.original_processor)

I can make a PR in the future but in a messy situation at the moment.