Corrected dtype setup

modeling_hymba.py

@@ -1709,9 +1709,9 @@ class HymbaBlock(nn.Module):
         ## Attention Head    
         if self.reuse_kv:
             assert kv_last_layer is not None
-            attn_outputs, attn_key_value = self.self_attn(attention_mask=attention_mask, position_ids=position_ids, query_states=query_states, kv_last_layer=kv_last_layer, use_swa=use_swa, use_cache=use_cache, past_key_value=cache_params, target_dtype=self.in_proj.weight.type)
+            attn_outputs, attn_key_value = self.self_attn(attention_mask=attention_mask, position_ids=position_ids, query_states=query_states, kv_last_layer=kv_last_layer, use_swa=use_swa, use_cache=use_cache, past_key_value=cache_params, target_dtype=self.in_proj.weight.dtype)
         else:
-            attn_outputs, attn_key_value = self.self_attn(attention_mask=attention_mask, position_ids=position_ids, query_states=query_states, key_states=key_states, value_states=value_states, use_swa=use_swa, use_cache=use_cache, past_key_value=cache_params, target_dtype=self.in_proj.weight.type)
+            attn_outputs, attn_key_value = self.self_attn(attention_mask=attention_mask, position_ids=position_ids, query_states=query_states, key_states=key_states, value_states=value_states, use_swa=use_swa, use_cache=use_cache, past_key_value=cache_params, target_dtype=self.in_proj.weight.dtype)
 
 
         if not self.pure_attn: