Don't return logits during training for PyTorch baseline

This improves perf somewhat, since currently it's always returning logits (which thus need to be materialized).

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):