[ROCm][Bugfix] Fix init-time bias dtype cast when gate.out_dtype is None (#41405)

Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
2026-05-02 00:13:15 -04:00
parent d58c42e19c
commit c293ccc58e
1 changed files with 2 additions and 1 deletions
@@ -351,8 +351,9 @@ class DeepseekV2MoE(nn.Module):
            self.is_rocm_aiter_moe_enabled
            and self.gate.e_score_correction_bias is not None
        ):
+            gate_out_dtype = self.gate.out_dtype or self.gate.weight.dtype
            self.gate.e_score_correction_bias.data = (
-                self.gate.e_score_correction_bias.data.to(self.gate.out_dtype)
+                self.gate.e_score_correction_bias.data.to(gate_out_dtype)
            )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: