From 3d901f4e5bd71669f37332f58097a4875b5388ef Mon Sep 17 00:00:00 2001 From: Ahmed Khaled Date: Wed, 3 Jun 2026 11:29:48 -0700 Subject: [PATCH] internal PiperOrigin-RevId: 926160988 --- init2winit/model_lib/rope_nanodo.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/init2winit/model_lib/rope_nanodo.py b/init2winit/model_lib/rope_nanodo.py index 803cfff9..bb2d7348 100644 --- a/init2winit/model_lib/rope_nanodo.py +++ b/init2winit/model_lib/rope_nanodo.py @@ -217,9 +217,9 @@ def __call__(self, x_BxLxD: jax.Array): q_BxLxHxDh, k_BxLxHxDh = apply_rope(q_BxLxHxDh, k_BxLxHxDh, self.freqs_cis) out_BxLxHxDh = jax.nn.dot_product_attention( - q_BxLxHxDh, - k_BxLxHxDh, - v_BxLxHxDh, + q_BxLxHxDh.astype(cfg.dtype), + k_BxLxHxDh.astype(cfg.dtype), + v_BxLxHxDh.astype(cfg.dtype), is_causal=cfg.is_causal, scale=1.0, )