Batch size greater than 1 results in IndexError

[Xyem]

It would appear that when this extension is enabled, I am unable to run generations where the batch size is greater than 1, giving an error like IndexError: list index out of range.

Is this an intentional/known limitation?

[Aisper]

Hey, check out my PR