* Fix kokoro batch issue * code * fix batch size --------- Co-authored-by: Freddy Boulton <freddyboulton@hf-freddy.local>