consistent shape between logits and generated_ids

Author: lapp0

exllamav2/generator/sampler.py

@@ -364,7 +364,9 @@ def sample(
         # Apply logits processor
 
         if settings.logits_processor:
-            generated_ids = sequence_ids[:, input_ids.shape[1]:]
+            generated_ids = sequence_ids[:, input_ids.shape[1]:].view(
+                logits.shape[:-1] + sequence_ids.shape[-1:]  # ensure consistent batch dimensions
+            )
             logits = settings.logits_processor(generated_ids, logits)
 
         # Prepare filter