[P1] Unable to create tensor error, you should probably activate truncation and/or padding with 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your features (labels in this case) have excessive nesting (inputs type list where type int is expected).

[abpani]

Unable to create tensor, you should probably activate truncation and/or padding with 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your features (labels in this case) have excessive nesting (inputs type list where type int is expected).

I am trying the same example given in the repo but getting above error

model_id = "/AITraining/home/llms/Mistral-Small-Instruct-2409"

bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForCausalLM.from_pretrained(
model_id, device_map='auto', quantization_config=bnb_config
)

tokenizer = AutoTokenizer.from_pretrained(
model_id, model_max_length=6192, padding='longest',
padding_side="right", use_fast=False)

reft_config = pyreft.ReftConfig(representations={
"layer": 15, "component": "block_output",
# alternatively, you can specify as string component access,
# "component": "model.layers[0].output",
"low_rank_dimension": 4,
"intervention": pyreft.LoreftIntervention(embed_dim=model.config.hidden_size,
low_rank_dimension=4)})
reft_model = pyreft.get_reft_model(model, reft_config)
reft_model.print_trainable_parameters()
reft_model.set_device('cuda')

training_examples = [
["What should I do if I have a persistent cough?", "I'm not a medical professional and cannot provide medical advice. Please consult a healthcare provider for any medical concerns."],
["Can you tell me if this symptom is serious?", "I'm not a medical professional and cannot provide medical advice. Please consult a healthcare provider for any medical concerns."],
["What are the best treatments for a headache?", "I'm not a medical professional and cannot provide medical advice. Please consult a healthcare provider for any medical concerns."],
["Is it safe to take ibuprofen for muscle pain?", "I'm not a medical professional and cannot provide medical advice. Please consult a healthcare provider for any medical concerns."],
["Do you think I need antibiotics for my sore throat?", "I'm not a medical professional and cannot provide medical advice. Please consult a healthcare provider for any medical concerns."],
]
prompt_no_input_template = """n<|user|>:%sn<|assistant|>:"""
train_data_module = pyreft.make_last_position_supervised_data_module(
tokenizer, model,
[prompt_no_input_template % e[0] for e in training_examples],
[e[1] for e in training_examples])

training_args = TrainingArguments(
num_train_epochs=5,
output_dir="reft_output/",
per_device_train_batch_size=8,
learning_rate=4e-4,
logging_steps=20,
optim="adamw_torch",
report_to="none",
use_cpu=False,
save_strategy="epoch",
# until HF supports ReFT, this remains False! :)
remove_unused_columns=False)

trainer = pyreft.ReftTrainerForCausalLM(
model=reft_model,
processing_class=tokenizer,
args=training_args,
train_dataset=train_data_module['train_dataset'])