From ad55575326443db6deda6e19126ebf136c66d8b2 Mon Sep 17 00:00:00 2001
From: George Shakan <43767775+georgeshakan@users.noreply.github.com>
Date: Wed, 18 Feb 2026 10:42:11 -0500
Subject: [PATCH] Fix bug in setting precision (#538)

---
 nanochat/common.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nanochat/common.py b/nanochat/common.py
index 9bcd5dd..2dd0792 100644
--- a/nanochat/common.py
+++ b/nanochat/common.py
@@ -170,7 +170,7 @@ def compute_init(device_type="cuda"): # cuda|cpu|mps
 
     # Precision
     if device_type == "cuda":
-        torch.backends.fp32_precision = "tf32" # uses tf32 instead of fp32 for matmuls
+        torch.set_float32_matmul_precision("high") # uses tf32 instead of fp32 for matmuls, see https://docs.pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html
 
     # Distributed setup: Distributed Data Parallel (DDP), optional, and requires CUDA
     is_ddp_requested, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()