diff --git a/train_gpt2.py b/train_gpt2.py
index f2fa68c..1446e61 100644
--- a/train_gpt2.py
+++ b/train_gpt2.py
@@ -147,12 +147,11 @@ class GPT(nn.Module):
             # if we are given some desired targets also calculate the loss
             logits = self.lm_head(x)
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+            return None, loss
         else:
             # inference-time mini-optimization: only forward the lm_head on the very last position
             logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
-            loss = None
-
-        return logits, loss
+            return logits, None
 
     @classmethod
     def from_pretrained(cls, model_type):