AnglE: RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

I am trying to train my model using LLAMA-v2-nli. I was able to do so with the bert-nli model but when I try to run with LLAMA I get the following error:

RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

`from angle_emb import AnglE, AngleDataTokenizer angle = AnglE.from_pretrained(‘NousResearch/Llama-2-7b-hf’, pretrained_lora_path=‘SeanLee97/angle-llama-7b-nli-v2’).cuda() train_ds = ds[‘train’].shuffle().map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8) valid_ds = ds[‘valid’].map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8) test_ds = ds[‘test’].map(AngleDataTokenizer(angle.tokenizer, angle.max_length), num_proc=8)

angle.fit( train_ds=train_ds, valid_ds=test_ds, output_dir=‘ckpts/sts-b’, batch_size=16, epochs=5, learning_rate=2e-5, save_steps=100, eval_steps=1000, warmup_steps=0, gradient_accumulation_steps=1, loss_kwargs={ ‘w1’: 1.0, ‘w2’: 1.0, ‘w3’: 1.0, ‘cosine_tau’: 20, ‘ibn_tau’: 20, ‘angle_tau’: 1.0 }, fp16=True, logging_steps=100 )`

I used the same code for bert (loaded the bert model instead) and it works no issues