Unhook cross attention processor (#53)

castorini · Nov 30, 2023 · b32e133 · b32e133
1 parent a323129
commit b32e133
Showing 1 changed file with 4 additions and 0 deletions.
diff --git 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())))