Method post_forward_hook

knn_utils/saveEmbedMulti.py:415–420

(self, module, input, output)

        return self.original_forward_func(input_ids=input_ids, labels=labels, attention_mask=attention_mask, **kwargs)

    def post_forward_hook(self, module, input, output):
        captured_keys = self.activation_capturer.captured

        self._save_step_data(keys = captured_keys, vals = self.labels)
        
        return output

    def register_hook(self, layer, func, pre=False):
        handle = layer.register_forward_pre_hook(func) if pre else layer.register_forward_hook(func)