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

cleanup

Browse Source
This commit is contained in:
Salman Mohammadi 2025-11-03 12:28:15 +00:00
parent e0e168dacd
commit e243767cc3
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
scripts/base_train.py
View File

@ -112,8 +112,7 @@ with torch.device("meta"):
model.to_empty(device=device) model.to_empty(device=device)
model.init_weights() model.init_weights()
orig_model = model # original, uncompiled model, for saving raw model state_dict orig_model = model # original, uncompiled model, for saving raw model state_dict
eval_model = model eval_model = torch.compile(model, fullgraph=True, dynamic=True) # eval model compiled for dynamic shapes
eval_model = torch.compile(eval_model, fullgraph=True, dynamic=True)
model = torch.compile(model, fullgraph=True, dynamic=False) # TODO: dynamic True/False think through model = torch.compile(model, fullgraph=True, dynamic=False) # TODO: dynamic True/False think through
num_params = sum(p.numel() for p in model.parameters()) num_params = sum(p.numel() for p in model.parameters())
print0(f"Number of parameters: {num_params:,}") print0(f"Number of parameters: {num_params:,}")