Update modeling_quiet.py

modeling_quiet.py

@@ -1119,7 +1119,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         max_new_tokens: Optional[int] = None,
         temperature: float = 1.1,
         **kwargs,
-    ):	
+    ):
         if isinstance(input_ids, str):
             input_ids = self.tokenizer(input_ids, return_tensors="pt").input_ids
 
@@ -1128,7 +1128,17 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
             attention_mask = torch.ones_like(input_ids)
 
         from .generate import generate
-        return generate(self, input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
+
+        output = generate(
+            self,
+            input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            **kwargs,
+        )
+
+        return output.sequences
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)