Method pre_forward_hook

knn_utils/saveEmbedMulti.py:150–153

(self, input_ids=None, attention_mask=None, labels=None, **kwargs)

        return dists, knns

    def pre_forward_hook(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        self.labels = labels
        self.input_ids = input_ids
        return self.original_forward_func(input_ids=input_ids, labels=labels, attention_mask=attention_mask, **kwargs)

    def post_forward_hook(self, module, input, output):
        batch, time_dim, vocab_size = output.shape