Update cognitivess_model/modeling_cognitivess.py

cognitivess_model/modeling_cognitivess.py

@@ -237,9 +237,9 @@ class CognitivessForQuestionAnswering(LlamaPreTrainedModel):
             end_loss = loss_fct(end_logits, end_positions)
             loss = (start_loss + end_loss) / 2
 
-            if not return_dict:
-                output = (start_logits, end_logits) + outputs[1:]
-                return ((loss,) + output) if loss is not None else output
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
 
         return QuestionAnsweringModelOutput(
             loss=loss,