Update modeling_quiet.py

modeling_quiet.py

@@ -158,7 +158,7 @@ class QuietRMSNorm(nn.Module):
 
     def forward(self, hidden_states):
         input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.bfloat16)
+        hidden_states = hidden_states.to(torch.float32)
         variance = hidden_states.pow(2).mean(-1, keepdim=True)
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         return hidden_states.to(input_dtype) * self.weight.to(hidden_states.device)
@@ -327,9 +327,12 @@ class QuietAttention(nn.Module):
                 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
             )
         bsz, q_len, _ = hidden_states.size()
-
+        
+        hidden_states = hidden_states.to(self.q_proj.weight.dtype)
         query_states = self.q_proj(hidden_states)
+        hidden_states = hidden_states.to(self.k_proj.weight.dtype) 
         key_states = self.k_proj(hidden_states)
+        hidden_states = hidden_states.to(self.v_proj.weight.dtype) 
         value_states = self.v_proj(hidden_states)
 
         query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)