this?

stanford-crfm · Dec 2, 2024 · 5052f36 · 5052f36
1 parent 4fff4c3
commit 5052f36
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_attention.py b/tests/test_attention.py
@@ -241,6 +241,8 @@ def test_segment_ids_are_respected(impl):
     keys = hax.named(keys, (KPos, Head))
     values = hax.named(values, (KPos, Head))
 
+    query, keys, values = jax.device_put([query, keys, values])
+
     segment_ids = np.array([0, 0, 0] + [1] * (L - 3), dtype=np.int32)
     segment_ids = jax.device_put(segment_ids)
     segment_ids = hax.named(segment_ids, (Pos,))