Update modelling_single.py

aehrc · Aug 22, 2024 · 9aea762 · 9aea762
1 parent 6969c23
commit 9aea762
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/modules/transformers/single_model/modelling_single.py b/modules/transformers/single_model/modelling_single.py
@@ -54,6 +54,7 @@ def forward(
         pixel_values: Optional[torch.Tensor] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
     ) -> Union[Tuple, ModelOutput]:
 
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
@@ -405,4 +406,4 @@ def split_and_decode_sections(self, token_ids, special_token_ids, tokenizer: Pre
 
                 sections[j].append(section_string)
 
-        return tuple(sections.values())
+        return tuple(sections.values())