From 69cc309967f6e6a71cf1c9efa717ff14ca74da87 Mon Sep 17 00:00:00 2001
From: spjosyula <spjosyula2005@outlook.com>
Date: Wed, 24 Dec 2025 01:56:45 +0530
Subject: [PATCH] perf: remove redundant dtype conversion in apply_rotary_emb

---
 nanochat/gpt.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/nanochat/gpt.py b/nanochat/gpt.py
index 9a80c7c..0503d4e 100644
--- a/nanochat/gpt.py
+++ b/nanochat/gpt.py
@@ -44,9 +44,7 @@ def apply_rotary_emb(x, cos, sin):
     x1, x2 = x[..., :d], x[..., d:] # split up last time into two halves
     y1 = x1 * cos + x2 * sin # rotate pairs of dims
     y2 = x1 * (-sin) + x2 * cos
-    out = torch.cat([y1, y2], 3) # re-assemble
-    out = out.to(x.dtype) # ensure input/output dtypes match
-    return out
+    return torch.cat([y1, y2], 3) # re-assemble
 
 class CausalSelfAttention(nn.Module):
     def __init__(self, config, layer_idx):