Runtime error on 2nd epoch, trying lora only

[wasifferoze]

I ran into this problem, I am playing with multirc data and using lora only. First epoch was successful and also validation epoch, on second epoch this error happened, so I am wondering why on 2nd epoch.

Traceback (most recent call last): File "/root/data1/LOMO/src/train_lomo_lora.py", line 181, in <module> train() File "/root/data1/LOMO/src/train_lomo_lora.py", line 174, in train trainer.train() File "/root/data1/LOMO/src/lomo_lora_trainer.py", line 202, in train loss = get_loss(outs.logits, batch['labels'], self.training_args.clip_loss_value) File "/root/data1/LOMO/src/utils.py", line 381, in get_loss loss = loss_fct(shift_logits.view(shift_labels.shape[0] * shift_labels.shape[1], -1), RuntimeError: shape '[4096, -1]' is invalid for input of size 117504000