Update modeling_quiet.py

modeling_quiet.py

@@ -774,7 +774,7 @@ class QuietSdpaAttention(QuietAttention):
 				raise ValueError(
 					f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
 				)
-			attention_mask = attention_mask.to(query_states.dtype)
+			
 		# SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
 		# Reference: https://github.com/pytorch/pytorch/issues/112577.
 		if query_states.device.type == "cuda" and attention_mask is not None:
@@ -1674,15 +1674,12 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 					base_attention_mask = base_attention_mask.view(1, 1, seq_len, seq_len)
 					base_attention_mask = base_attention_mask.repeat(input_ids.shape[0], 1, 1, 1)
 					attention_mask = base_attention_mask
-					# breakpoint()
 				elif attention_mask.dim() == 2:
 					if seq_len + past_key_values_length != attention_mask.shape[-1]:
-						# breakpoint()
 						attention_mask = torch.cat(
 							[torch.ones((attention_mask.shape[0], past_key_values_length), dtype=attention_mask.dtype, device=attention_mask.device), attention_mask],
 							dim=-1
 						)
-					# # if the attention mask 
 					attention_mask = _prepare_4d_causal_attention_mask(
 						attention_mask,
 						(batch_size, seq_len),
@@ -1700,10 +1697,8 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 				use_cache=use_cache,
 				output_attentions=output_attentions,
 				output_hidden_states=output_hidden_states,
-				# output_router_logits=output_router_logits,
 				return_dict=return_dict,
 			)
-
 			prev_hidden_states = hidden_states
 			hidden_states = outputs[0]
 			prev_rm_logits = rm_logits  # for policy gradient