fixes

Lightning-AI · Oct 2, 2024 · 97406e3 · 97406e3
1 parent b11fcdb
commit 97406e3
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/litgpt/model.py b/litgpt/model.py
@@ -461,10 +461,9 @@ def build_rope_cache(
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: Cosine and sine caches for RoPE.
     """
-    assert n_elem % 2 == 0, "n_elem (head dimension) must be even"
 
-    # Compute the initial inverse frequencies (theta)
-    theta = 1.0 / (base ** (torch.arange(0, n_elem // 2, device=device).float() / (n_elem // 2)))
+    # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$	    assert n_elem % 2 == 0, "n_elem (head dimension) must be even"
+    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem))
 
     if extra_config is not None:
         # Extract configuration parameters