From c0cc99cdadb8ecc2003ddd3233f8a04e0d6531ff Mon Sep 17 00:00:00 2001
From: Charles Martin <charlesmartin14@gmail.com>
Date: Mon, 24 Aug 2026 13:35:47 -0700
Subject: [PATCH] Fix PyTorch/XLA evaluation inference-mode crash

---
 .../nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py
index a5adc95..71aeb63 100644
--- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py
+++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/evaluation.py
@@ -128,7 +128,7 @@ def fixed_bleu_probe(
     )
 
 
-@torch.inference_mode()
+@torch.no_grad()
 def evaluate_probe(
     model: nn.Module,
     probe: Iterable[tuple[torch.Tensor, torch.Tensor]],
@@ -232,7 +232,7 @@ def _cpu_bleu_model(model) -> nn.Module:
     return cpu_model
 
 
-@torch.inference_mode()
+@torch.no_grad()
 def evaluate_bleu(
     model,
     probe: BleuProbe,