Loss calculation should not permanently change shapes of logits and targets

bigram.py

@@ -74,9 +74,7 @@ def forward(self, idx, targets=None):
             loss = None
         else:
             B, T, C = logits.shape
-            logits = logits.view(B*T, C)
-            targets = targets.view(B*T)
-            loss = F.cross_entropy(logits, targets)
+            loss = F.cross_entropy(logits.view(B*T, C), targets.view(B*T))
 
         return logits, loss
 

gpt.py

@@ -103,7 +103,7 @@ def forward(self, x):
         out = self.dropout(self.proj(out))
         return out
 
-class FeedFoward(nn.Module):
+class FeedForward(nn.Module):
     """ a simple linear layer followed by a non-linearity """
 
     def __init__(self, n_embd):
@@ -126,7 +126,7 @@ def __init__(self, n_embd, n_head):
         super().__init__()
         head_size = n_embd // n_head
         self.sa = MultiHeadAttention(n_head, head_size)
-        self.ffwd = FeedFoward(n_embd)
+        self.ffwd = FeedForward(n_embd)
         self.ln1 = nn.LayerNorm(n_embd)
         self.ln2 = nn.LayerNorm(n_embd)
 
@@ -172,9 +172,7 @@ def forward(self, idx, targets=None):
             loss = None
         else:
             B, T, C = logits.shape
-            logits = logits.view(B*T, C)
-            targets = targets.view(B*T)
-            loss = F.cross_entropy(logits, targets)
+            loss = F.cross_entropy(logits.view(B*T, C), targets.view(B*T))
 
         return logits, loss