Compute_metrics caused training stopped during evalauation
Hi, I am training Mbart50 for conditional text generation. I prepared the dataset in the required format using datasets.load_dataset. Later, I loaded the model and tokenizer given asfrom transformers import MBartForConditionalGeneration, MBart50Tokenizer
from transformers import Trainer, TrainingArguments
tokenizer = MBart50Tokenizer.from_pretrained("facebook/mbart-large-50") 
model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50")I have specified the training arguments  astraining_args = TrainingArguments(
      output_dir=output_dir,           
      overwrite_output_dir = True,
      #do_train = True,
      num_train_epochs=1,           
      per_device_train_batch_size=4,  
      per_device_eval_batch_size=4,    
      save_steps=100,                  
      save_total_limit=3,              
      evaluation_strategy='steps',     
      save_strategy = 'steps',
      learning_rate = 8e-5,             
      lr_scheduler_type = 'linear',      
      weight_decay=0.01,             
      logging_dir='./logs',          
      logging_steps=10,
)and trainer astrainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
)I am training usingtrainer.train()and training is done fine.But later I wanted to include  compute_metrics in the tarining. I have defined the compute_metrics asfrom datasets import load_metric
from transformers import EvalPrediction
rouge = load_metric('rouge')
def compute_metrics(p: EvalPrediction):
  predictions = p.predictions
  labels= p.label_ids
  decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
  labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
  decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
  decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
  decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]
  result = rouge.compute(decoded_preds, decoded_labels)
  result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
  prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
  results["gen_len"] = np.mean(prediction_lens)
  results = {k: round(v, 4) for k, v in result.items()}
  return resultsand then modified the trainer astrainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
    compute_metrics = compute_metrics
)When I run the training usingtrainer.train(), it trains till first logging step and starts doing evaluation but in between the evaluation it stops.I am not able to understand what is going wrong here. Can somebody help me understand and fix it