fix: flatten noncontiguous AttnRes targets

experiments/k3/attnres/train.py

@@ -390,7 +390,9 @@ def learning_rate(step: int, total_steps: int) -> float:
 
 
 def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
-    return F.cross_entropy(logits.float().view(-1, VOCABULARY), targets.view(-1))
+    return F.cross_entropy(
+        logits.float().reshape(-1, VOCABULARY), targets.reshape(-1)
+    )
 
 
 @torch.no_grad()
@@ -410,8 +412,8 @@ def evaluate(
         with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
             logits, _ = model(inputs)
         loss = F.cross_entropy(
-            logits.float().view(-1, VOCABULARY),
-            targets.view(-1),
+            logits.float().reshape(-1, VOCABULARY),
+            targets.reshape(-1),
             reduction="sum",
         )
         loss_sum += loss.detach().cpu().item()