Update modeling_quiet.py

modeling_quiet.py

@@ -1412,7 +1412,17 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
                         past_key_values_length,
                         sliding_window=self.config.sliding_window,
                     )
-
+            if attention_mask is not None:
+                if attention_mask.dim() == 2:
+                    # Expand the attention mask to have dimensions (batch_size, 1, 1, seq_length)
+                    attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+                elif attention_mask.dim() == 3:
+                    # Expand the attention mask to have dimensions (batch_size, 1, seq_length, seq_length)
+                    attention_mask = attention_mask.unsqueeze(1)
+                else:
+                    raise ValueError(
+                        f"Attention mask should have 2 or 3 dimensions, but got {attention_mask.dim()} dimensions."
+                    )
             outputs = self.model(
                 # input_ids=input_ids,
                 attention_mask=attention_mask,