tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-01-21 02:44:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

slightly nicer error message

Browse Source
This commit is contained in:
Matěj Kripner 2025-12-09 12:46:48 +01:00
parent f1bf69d562
commit bbc57da7d5
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/adamw.py
View File

@ -27,7 +27,7 @@ class DistAdamW(torch.optim.Optimizer):
for group in self.param_groups:
params: list[Tensor] = group["params"]
for base_i in range(len(params)):
assert params[base_i].shape[0] % world_size == 0, f"Parameter shape {params[base_i].shape} must be divisible by world size {world_size}"
assert params[base_i].shape[0] % world_size == 0, f"First dim of parameter shape {params[base_i].shape} must be divisible by world size {world_size}"
grad = params[base_i].grad
rank_size = grad.shape[0] // world_size
grad_slice = torch.empty_like(grad[:rank_size])