torch error while running the script ```llama-3.1-8b-32k-sampling.ipynb```

While running the example given here : https://github.com/aws-neuron/aws-neuron-samples/blob/master/torch-neuronx/transformers-neuronx/inference/llama-3.1-8b-32k-sampling.ipynb , I am getting the following error when running the line

neuron_model.sample(input_ids, sequence_length=32768, top_k=10)

Error :

Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
  File "/opt/aws_neuronx_venv_transformers_neuronx/lib/python3.10/site-packages/transformers_neuronx/llama/model.py", line 240, in sample
    result = sampling.sample_llama(
  File "/opt/aws_neuronx_venv_transformers_neuronx/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/aws_neuronx_venv_transformers_neuronx/lib/python3.10/site-packages/transformers_neuronx/sampling.py", line 371, in sample_llama
    return sample_loop_llama(
  File "/opt/aws_neuronx_venv_transformers_neuronx/lib/python3.10/site-packages/transformers_neuronx/sampling.py", line 331, in sample_loop_llama
    done_flags = torch.logical_or(done_flags, inputs == eos_token_id)
TypeError: logical_or(): argument 'other' (position 2) must be Tensor, not bool

My instance :

inf2.48xlarge
ami_image : Deep Learning AMI Neuron (Ubuntu 22.04)
ami_id : ami-05d498302130f9036

Followed this doc : https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/neuron-setup/multiframework/multi-framework-ubuntu22-neuron-dlami.html and run the command source /opt/aws_neuronx_venv_transformers_neuronx/bin/activate

Current Installed Libraries :

PyDispatcher==2.0.7
Pygments==2.18.0
pylint==3.2.6
pyOpenSSL==24.2.1
pyparsing==3.1.2
pyproject_hooks==1.1.0
PySocks==1.7.1
pytest==8.3.2
python-daemon==3.0.1
python-dateutil==2.9.0.post0
python-json-logger==2.0.7
pytz==2024.1
PyYAML==6.0.2
pyzmq==26.1.0
qtconsole==5.5.2
QtPy==2.4.1
queuelib==1.7.0
referencing==0.35.1
regex==2024.7.24
requests==2.31.0
requests-file==2.1.0
requests-toolbelt==1.0.0
requests-unixsocket==0.3.0
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rich==13.7.1
rpds-py==0.20.0
rsa==4.7.2
ruamel.yaml==0.18.6
ruamel.yaml.clib==0.2.8
s3transfer==0.10.2
safetensors==0.4.4
safety==2.3.5
scikit-learn==1.5.1
scipy==1.11.2
Scrapy==2.11.2

aws-neuron / aws-neuron-samples

torch error while running the script ```llama-3.1-8b-32k-sampling.ipynb``` #88