[HotFix]Fix dtype mismatch in nsa indexer on AMD device (#17518)

2026-01-22 12:17:02 +08:00
parent d6e2b88288
commit 3373545b9f
1 changed files with 1 additions and 1 deletions
@@ -186,7 +186,7 @@ class Indexer(MultiPlatformOp):
            self.hidden_size,
            self.n_heads,
            bias=False,
-            params_dtype=torch.bfloat16,
+            params_dtype=torch.bfloat16 if _is_cuda else torch.float32,
            prefix=add_prefix("weights_proj", prefix),
        )
        self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)