Encountered this error while training, how to solve this TypeError: ViTModel.forward() got an unexpected keyword argument 'num_items_in_batch'