LIT: Fix assertion error when generation stops short of max_length

PiperOrigin-RevId: 681096247

lit_nlp/examples/prompt_debugging/transformers_lms.py

@@ -279,12 +279,12 @@ def _get_batched_outputs(
     if self.framework == MLFramework.PT:
       encoded_inputs = encoded_inputs.to(self.device)
 
-    outputs = self.model.generate(
-        encoded_inputs["input_ids"],
-        attention_mask=encoded_inputs["attention_mask"],
-        max_length=self.max_length,
-    )
-    ntok_out = self.max_length - encoded_inputs["input_ids"].shape[1]
+    outputs = self.model.generate(**encoded_inputs, max_length=self.max_length)
+
+    if isinstance(outputs, transformers.utils.ModelOutput):
+      outputs = outputs.sequences
+
+    ntok_out = outputs.shape[1] - encoded_inputs["input_ids"].shape[1]
 
     responses = self.tokenizer.batch_decode(
         outputs[:, -ntok_out:], skip_special_tokens=True