From ed519b0f24ea1620a81653012103122a2e367d83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ph=C3=BAc=20H=2E=20L=C3=AA=20Kh=E1=BA=AFc?= Date: Fri, 17 Oct 2025 17:21:25 +0700 Subject: [PATCH] Update engine.py with correct error message on assert --- nanochat/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nanochat/engine.py b/nanochat/engine.py index de1253a..eb3fcac 100644 --- a/nanochat/engine.py +++ b/nanochat/engine.py @@ -83,7 +83,7 @@ class KVCache: for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)): if ix in [0, 1, 3, 5]: # num_layers, batch_size, num_heads, head_dim must match - assert dim1 == dim2, f"Batch dim mismatch: {dim1} != {dim2}" + assert dim1 == dim2, f"Dim {ix} mismatch: {dim1} != {dim2}" elif ix == 2: # batch_size can be expanded assert dim1 == dim2 or dim2 == 1, f"Batch dim mismatch: {dim1} != {dim2}"