remove redundant type convert

Author: HaochenYuan

megatron/core/transformer/moe/moe_utils.py

@@ -116,7 +116,7 @@ def switch_load_balancing_loss_func(
     # Apply padding mask to probs if provided
     if padding_mask is not None:
         # padding_mask: [num_tokens], probs: [num_tokens, num_experts]
-        mask_expanded = padding_mask.unsqueeze(-1).to(probs.dtype)
+        mask_expanded = padding_mask.unsqueeze(-1)
         probs = probs * mask_expanded
 
     aggregated_probs_per_expert = probs.sum(dim=0)
@@ -145,7 +145,7 @@ def z_loss_func(logits, z_loss_coeff, padding_mask: Optional[torch.Tensor] = Non
 
     if padding_mask is not None:
         # Only compute z_loss for non-padding tokens
-        z_loss_values = z_loss_values * padding_mask.to(z_loss_values.dtype)
+        z_loss_values = z_loss_values * padding_mask
         # Compute mean over valid tokens only
         num_valid_tokens = padding_mask.sum()
         z_loss = z_loss_values.sum() / torch.clamp(num_valid_tokens, min=1.0) * z_loss_coeff