Update modeling_quiet.py

modeling_quiet.py

@@ -1233,7 +1233,10 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         #         )
         #         if labels is not None:
         #             loss += self.router_aux_loss_coef * aux_loss.to(loss.device)
-                    
+        if input_ids.dim() == 1:
+            input_ids = input_ids.unsqueeze(0)
+            attention_mask = attention_mask.unsqueeze(0) if attention_mask is not None else None
+            
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states