From 75972ed67e4e327085e2f384832bf2d04fde6a5b Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Wed, 15 May 2024 15:36:00 +0200 Subject: [PATCH 1/2] Add inference mode to ExtractiveReader --- haystack/components/readers/extractive.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/haystack/components/readers/extractive.py b/haystack/components/readers/extractive.py index 573aa53038..62e1ce0022 100644 --- a/haystack/components/readers/extractive.py +++ b/haystack/components/readers/extractive.py @@ -602,7 +602,8 @@ def run( cur_input_ids = input_ids[start_index:end_index] cur_attention_mask = attention_mask[start_index:end_index] - output = self.model(input_ids=cur_input_ids, attention_mask=cur_attention_mask) + with torch.inference_mode(): + output = self.model(input_ids=cur_input_ids, attention_mask=cur_attention_mask) cur_start_logits = output.start_logits cur_end_logits = output.end_logits if num_batches != 1: From 35864f12f448fd72f1611943782f12449f2a636c Mon Sep 17 00:00:00 2001 From: Sebastian Husch Lee Date: Wed, 15 May 2024 15:38:00 +0200 Subject: [PATCH 2/2] Add release notes --- releasenotes/notes/add-inf-mode-reader-e6eb79920e73c956.yaml | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 releasenotes/notes/add-inf-mode-reader-e6eb79920e73c956.yaml diff --git a/releasenotes/notes/add-inf-mode-reader-e6eb79920e73c956.yaml b/releasenotes/notes/add-inf-mode-reader-e6eb79920e73c956.yaml new file mode 100644 index 0000000000..4dbe549f13 --- /dev/null +++ b/releasenotes/notes/add-inf-mode-reader-e6eb79920e73c956.yaml @@ -0,0 +1,4 @@ +--- +enhancements: + - | + Adds inference mode to model call of the ExtractiveReader. This prevents gradients from being calculated during inference time in pytorch.