jbk

modelLM.py

@@ -47,10 +47,8 @@ class OBILanguageModel(PreTrainedModel):
         x = self.ln2(x)
         logits = self.lm_head(x)
 
-        if targets is None:
-            loss = None
-        else:
-            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))
+            # Always compute the loss, and set it to None if targets are not provided
+        loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1)) if targets is not None else None
 
         return (logits, loss)