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

fix: cast bf16 to fp32 on MPS (like CPU) to avoid dtype issues

Browse Source
This commit is contained in:
Dipesh Babu 2025-11-03 16:00:56 -05:00
parent a83646e098
commit 7a40ee77b4
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/checkpoint_manager.py
View File

@ -65,7 +65,7 @@ def build_model(checkpoint_dir, step, device, phase):
""" """
assert phase in ["train", "eval"], f"Invalid phase: {phase}" assert phase in ["train", "eval"], f"Invalid phase: {phase}"
model_data, optimizer_data, meta_data = load_checkpoint(checkpoint_dir, step, device, load_optimizer=False) model_data, optimizer_data, meta_data = load_checkpoint(checkpoint_dir, step, device, load_optimizer=False)
if device.type == "cpu": if device.type in {"cpu", "mps"}:
# Convert bfloat16 tensors to float for CPU inference # Convert bfloat16 tensors to float for CPU inference
model_data = { model_data = {
k: v.float() if v.dtype == torch.bfloat16 else v k: v.float() if v.dtype == torch.bfloat16 else v