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