From b4fee701511613142c2b09f3da13e572961dac9e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 06:36:22 +0000 Subject: [PATCH] Fix device mismatch in test --- tests/test_o_projection_b200.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_o_projection_b200.py b/tests/test_o_projection_b200.py index edf07e83..1640db0d 100644 --- 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 ---")