Reproduction

The prompt_ids.size(1) maybe 0 when completion_ids.size(1) >= self.max_length.

fixed codes:

prompt_length = prompt_ids.size(1)
if prompt_length > 0:
    logits = output.logits[:, prompt_length - 1 : -1]  
else:
    logits = output.logits

System Info

nothing

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete