Pick a slot from 0 to batch_size-1 during run_interactive.py

[bhavya01]

To fix this index error
warnings.warn("Some donated buffers were not usable:"
---- Streaming decode started on #slot32.
new_pos Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>
cache_seq_len 2048
Traceback (most recent call last):
File "/home/bbahl/jetstream-pytorch/run_interactive.py", line 168, in
app.run(main)
File "/home/bbahl/.local/lib/python3.10/site-packages/absl/app.py", line 308, in run
_run_main(main, args)
File "/home/bbahl/.local/lib/python3.10/site-packages/absl/app.py", line 254, in _run_main
sys.exit(main(argv))
File "/home/bbahl/jetstream-pytorch/run_interactive.py", line 146, in main
token_id = slot_tokens[slot, 0].item()
File "/home/bbahl/.local/lib/python3.10/site-packages/jax/_src/array.py", line 348, in getitem
return lax_numpy._rewriting_take(self, idx)
File "/home/bbahl/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 4604, in _rewriting_take
return _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,
File "/home/bbahl/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 4613, in _gather
indexer = _index_to_gather(shape(arr), idx) # shared with _scatter_update
File "/home/bbahl/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 4831, in _index_to_gather
raise IndexError(f"index is out of bounds for axis {x_axis} with size 0")
IndexError: index is out of bounds for axis 0 with size 0

[FanhaiLu1]

Thanks for fixing it! I submitted the fit in multiple host but forgot to change this file.