From 4a332039f1937efa9828c262db75554741a2c237 Mon Sep 17 00:00:00 2001
From: Matthias Seeger <matthis@amazon.de>
Date: Tue, 30 Sep 2025 14:21:04 +0200
Subject: [PATCH] Fix to avoid RuntimeError: Failed to create function state
 object for: div_true_strided_bfloat_long

---
 litgpt/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/litgpt/model.py b/litgpt/model.py
index 01ea83ad4a..0321f3fad0 100644
--- a/litgpt/model.py
+++ b/litgpt/model.py
@@ -806,7 +806,7 @@ def build_rope_cache(
             theta = theta / factor
 
     # Create position indices `[0, 1, ..., seq_len - 1]`
-    seq_idx = torch.arange(seq_len, device=device) / condense_ratio
+    seq_idx = torch.arange(seq_len, device=device).float() / condense_ratio
 
     # Calculate the product of position index and $\theta_i$
     idx_theta = torch.outer(seq_idx, theta).repeat(1, 2)