tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-06 04:12:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

tiny fix to comment

Browse Source 
Update engine.py with correct error message on assert
This commit is contained in:
Andrej 2025-11-01 07:43:57 -07:00 committed by GitHub
commit ad39db5a23
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
nanochat/engine.py
View File

@ -109,7 +109,7 @@ class KVCache:
for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)): for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)):
if ix in [0, 1, 3, 5]: if ix in [0, 1, 3, 5]:
# num_layers, batch_size, num_heads, head_dim must match # 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: elif ix == 2:
# batch_size can be expanded # batch_size can be expanded
assert dim1 == dim2 or dim2 == 1, f"Batch dim mismatch: {dim1} != {dim2}" assert dim1 == dim2 or dim2 == 1, f"Batch dim mismatch: {dim1} != {dim2}"