Skip to content

Commit 75972ed

Browse files
committed
Add inference mode to ExtractiveReader
1 parent c8d53b3 commit 75972ed

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

Diff for: haystack/components/readers/extractive.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -602,7 +602,8 @@ def run(
602602
cur_input_ids = input_ids[start_index:end_index]
603603
cur_attention_mask = attention_mask[start_index:end_index]
604604

605-
output = self.model(input_ids=cur_input_ids, attention_mask=cur_attention_mask)
605+
with torch.inference_mode():
606+
output = self.model(input_ids=cur_input_ids, attention_mask=cur_attention_mask)
606607
cur_start_logits = output.start_logits
607608
cur_end_logits = output.end_logits
608609
if num_batches != 1:

0 commit comments

Comments
 (0)