From b32e13310721a860bb21e74d409a210d4ded9112 Mon Sep 17 00:00:00 2001 From: Dave Lage Date: Wed, 29 Nov 2023 19:08:57 -0500 Subject: [PATCH] Unhook cross attention processor (#53) --- daam/trace.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/daam/trace.py b/daam/trace.py index 6932a4a..54b9f71 100644 --- a/daam/trace.py +++ b/daam/trace.py @@ -279,8 +279,12 @@ def __call__( return hidden_states 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) + @property def num_heat_maps(self): return len(next(iter(self.heat_maps.values())))