Upload modeling_limon.py

modeling_limon.py

@@ -86,8 +86,15 @@ class LimonFlowV1Model(PreTrainedModel):
         x = self.embeddings(input_ids) + self.pos_embeddings(pos)
         x = self.ode_solver(x)
         logits = self.head(x)
-        loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1)) if labels is not None else None
-        return {"loss": loss, "logits": logits} if loss is not None else logits
+        
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        return {
+            "logits": logits,
+            "loss": loss
+        }
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
         return {"input_ids": input_ids}