When I try to do the tensor_parallel on NLLB from meta, there is an error

[342215448]

Here is my code:
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import tensor_parallel as tp

tokenizer = AutoTokenizer.from_pretrained("/workspace/projects/nllb/nllb3.3b", use_auth_token=True, src_lang="deu_Latn")
model = AutoModelForSeq2SeqLM.from_pretrained("/workspace/projects/nllb/nllb3.3b", use_auth_token=True)
model = tp.tensor_parallel(model, ["cuda:0", "cuda:1"])

parallelize(model, num_gpus=2, fp16=True, verbose='detail')

model.to('cuda')

article = "Die Ware hat unter 20 Euro gekostet."
inputs = tokenizer(article, return_tensors="pt").to("cuda:0")

output1=translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["zho_Hans"], max_length=500)

output1=translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["zho_Hans"], max_length=500)

output1=translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["zho_Hans"])

out=tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]

print(out)

#############################################################
And here is the error:
Using automatic config: tensor parallel config not provided and no custom config registered for the model
The following patterns in state_rules were unused: ["re.compile('^model.decoder.embed_tokens.weight$')", "re.compile('^model.encoder.embed_tokens.weight$')"]
The following patterns in state_rules were unused: ["re.compile('^model.decoder.embed_tokens.weight$')", "re.compile('^model.encoder.embed_tokens.weight$')"]
Using ZeRO-3 sharding for 499712 non tensor-parallel parameters
Traceback (most recent call last):
File "test_nllb.py", line 14, in
output1=translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id["zho_Hans"], max_length=500)
File "/usr/local/lib/python3.8/dist-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/workspace/tools/transformers/src/transformers/generation/utils.py", line 1518, in generate
return self.greedy_search(
File "/workspace/tools/transformers/src/transformers/generation/utils.py", line 2335, in greedy_search
outputs = self(
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/tensor_parallel/pretrained_model.py", line 78, in forward
return self.wrapped_model(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/tensor_parallel/sharding.py", line 95, in forward
return self.module(*args, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/tensor_parallel/tensor_parallel.py", line 130, in forward
return parallel_apply(self.module_shards, inputs, kwargs_tup, self.devices)[self.output_device_index]
File "/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/parallel_apply.py", line 86, in parallel_apply
output.reraise()
File "/usr/local/lib/python3.8/dist-packages/torch/_utils.py", line 463, in reraise
raise exception
KeyError: Caught KeyError in replica 0 on device 0.
Original Traceback (most recent call last):
File "/usr/local/lib/python3.8/dist-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
output = module(*input, **kwargs)
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/workspace/tools/transformers/src/transformers/models/m2m_100/modeling_m2m_100.py", line 1335, in forward
outputs = self.model(
File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1110, in _call_impl
return forward_call(*input, **kwargs)
File "/workspace/tools/transformers/src/transformers/models/m2m_100/modeling_m2m_100.py", line 1220, in forward
last_hidden_state=encoder_outputs[0],
KeyError: 0

#############################################################
Do any guys know why😂😂😂😂