fix: NVBug 5385576 py_batch_idx issue#6153

hchings merged 2 commits intoNVIDIA:mainNVIDIA/TensorRT-LLM:mainfrom

hchings:nvbug_5385576hchings/TensorRT-LLM:nvbug_5385576Copy head branch name to clipboard