From 3c4be194e960f19ae679cc062aa57d6ae7f042ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mat=C4=9Bj=20Kripner?= Date: Sat, 3 Jan 2026 10:59:14 +0100 Subject: [PATCH] Add comment with logits shape --- nanochat/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nanochat/engine.py b/nanochat/engine.py index cc207e8..ef2867a 100644 --- a/nanochat/engine.py +++ b/nanochat/engine.py @@ -359,7 +359,7 @@ class Engine: ) else: logits = self.model.forward(ids, kv_cache=kv_cache_decode) - logits = logits[:, -1, :] + logits = logits[:, -1, :] # (B, vocab_size) def generate_batch(self, tokens, num_samples=1, **kwargs): """