I am encountering a ValueError while training my model using the Hugging Face Transformers library. The error message states that the model did not return a loss from the inputs, only the following keys: start_logits, end_logits. However, the inputs it received are input_ids, token_type_ids, attention_mask.

# Load smaller model
model = BertForQuestionAnswering.from_pretrained('bert-base-uncased')

# Define training arguments (using eval_strategy instead of evaluation_strategy)
training_args = TrainingArguments(
  output_dir='./results',
  eval_strategy='epoch',  # Update for compatibility
  learning_rate=2e-5,
  per_device_train_batch_size=4,  # Further reduce batch size
  per_device_eval_batch_size=4,  # Further reduce batch size
  num_train_epochs=3,
  weight_decay=0.01,
  save_steps=100,                  # Save model less frequently to reduce memory usage
  save_total_limit=1               # Limit the number of saved checkpoints to reduce memory usage
)

def compute_loss(model, inputs):
  start_logits, end_logits = model(**inputs)
  labels = inputs["start_positions"]  # Assuming "start_positions" is your ground truth label key

  # MLM loss calculation (replace with your preferred loss function)
  loss = F.cross_entropy(start_logits, labels, reduction="mean")

  return loss

# Trainer
trainer = Trainer(
  model=model,
  args=training_args,
  train_dataset=train_dataset,
  eval_dataset=train_dataset,  # Use train_dataset for simplicity, ideally you should split a validation set
  compute_metrics=compute_loss
)

# Train model
trainer.train()

How can I resolve this issue and ensure that the model returns the loss from the inputs during training?

I tried chatgpt and Gemini to get help and try to write code different way but didnt help.