Gibberish output for princeton-nlp/Sheared-LLaMA-1.3B with continuous batching

[pinak-p]

I'm seeing gibberish output with princeton-nlp/Sheared-LLaMA-1.3B when using continuous batching with transformers-neuronx and optimum neuron.

aws-neuronx-runtime-discovery 2.9
libneuronxla                  2.0.2335
neuronx-cc                    2.14.213.0+013d129b
optimum-neuron                0.0.24.dev0
torch-neuronx                 2.1.2.2.2.0

optimum-cli export neuron -m princeton-nlp/Sheared-LLaMA-1.3B --batch_size 4 --sequence_length 4096 --num_cores 2 --auto_cast_type fp16 ./neuron-model

https://github.com/huggingface/optimum-neuron/blob/aws_neuron_sdk_2.19/examples/text-generation/generation.py

['<s> One of my fondest memory is Ozi, Ozi Ozi\nPi (Greek: πι) is the symbol of a circle. A circle in geometry is called a "polytope", or "triple cover". In topology, a circle is a polytope of genus zero, or a cover of type one of genus zero. In topology, an object is called a "closed curve".\nA closed curve is called a "surface". Surfaces are defined in terms of π and ρ. The intersection πρ of the two curves at a point x is called a "point', '<s> One of my fondest memory is a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a a', '<s> One of my fondest memory is the F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F F', '<s> One of my fondest memory is the sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks sinks s']

[pinak-p]

@dacorvo

[dacorvo]

This used to work with SDK 2.18.0.