Error in NXD LLama Inference

Hi there, looks like we have a uncompleted function definition on this line. Can you take a look?

https://github.com/aws-neuron/neuronx-distributed/blob/4f954715f39b2cc9e628ded79274957401bea086/examples/inference/llama2/neuron_modeling_llama.py#L297 