Fix device mismatch in test

2026-05-19 06:36:22 +00:00
parent 6b4b9774d1
commit b4fee70151
1 changed files with 1 additions and 1 deletions
--- a/tests/test_o_projection_b200.py
+++ b/tests/test_o_projection_b200.py
@@ -245,7 +245,7 @@ def main():
    # (it's a plain BF16 tensor, not a quantized layer)

    # Build cos_sin_cache
-    cos_sin_cache = build_cos_sin_cache()
+    cos_sin_cache = build_cos_sin_cache().to(DEVICE)

    # Simulate attention output (what FlashMLA would produce)
    print("\n--- Simulating attention output ---")