Update generate.py

generate.py

@@ -184,7 +184,7 @@ def generate(
     self.rm_initialized = True
     self.original_mode = False
 
-    generated_token_ids, attention_mask = custom_generate(
+    output = custom_generate(
         self,
         input_ids=input_ids,
         attention_mask=attention_mask,
@@ -221,4 +221,4 @@ def generate(
         **model_kwargs,
     )
 
-    return generated_token_ids, attention_mask
+    return output