Method forward_wrapper

knn_utils/saveEmbedMulti.py:126–127

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

        # Create a wrapper that calls pre_forward_hook, ensuring that self still refers to KNNSaverMulti
        def forward_wrapper(input_ids=None, attention_mask=None, labels=None, **kwargs):
            return self.pre_forward_hook(input_ids=input_ids, attention_mask=attention_mask, labels=labels, **kwargs)

        # Override the model's forward with our wrapper
        model.forward = forward_wrapper