hmm

rankers/train/trainer.py

@@ -9,7 +9,7 @@
 from datasets import Dataset
 from transformers.trainer_utils import EvalLoopOutput, speed_metrics
 from transformers.integrations.deepspeed import deepspeed_init
-from dataclasses import dataclass
+from dataclasses import is_dataclass
 from .loss import LOSS_REGISTRY
 
 logger = logging.getLogger(__name__)
@@ -65,9 +65,9 @@ def compute_loss(
                 "The model did not return a loss from the inputs, only the following keys: "
                 f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
             )
-        # We don't use .loss here since the model may return tuples instead of ModelOutput.
-        loss = outputs.loss if isinstance(outputs, dataclass) else outputs[0]
-        
+
+        loss = outputs.loss if is_dataclass(outputs) else outputs[0]
+
         return (loss, outputs) if return_outputs else loss
 
     def compute_metrics(self, result_frame: pd.DataFrame):