Correct the output dtype of rmsnorm_func (#13)

- Correct the output dtype of rmsnorm_func (f2e665eb9ee6eae4abd08d7cb44fdf6422ee0c84)


Co-authored-by: ag0 <ag0@users.noreply.huggingface.co>

modeling_flash_llama.py

@@ -68,7 +68,7 @@ def rmsnorm_func(hidden_states, weight, variance_epsilon):
     hidden_states = hidden_states.to(torch.float32)
     variance = hidden_states.pow(2).mean(-1, keepdim=True)
     hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon)
-    return weight * hidden_states.to(input_dtype)
+    return (weight * hidden_states).to(input_dtype)
     
 
 class LlamaRMSNorm(nn.Module):