From 5f49906be6a859733ef7c6d323938f1b852f9b5e Mon Sep 17 00:00:00 2001
From: wuyang <5700876+banisherwy@user.noreply.gitee.com>
Date: Thu, 30 Jul 2026 06:39:29 +0800
Subject: [PATCH] fix: flatten noncontiguous AttnRes targets

---
 experiments/k3/attnres/train.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/experiments/k3/attnres/train.py b/experiments/k3/attnres/train.py
index 7f55078..b7c2583 100644
--- a/experiments/k3/attnres/train.py
+++ b/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()