Update generate.py

generate.py

@@ -207,4 +207,4 @@ def generate(
         **model_kwargs,
     )
 
-    return torch.tensor(generated_token_ids)
+    return generated_token_ids, attention_mask