Trainer.predict() does not return values in optuna search

I am trying to train some models on question answering mostly following this course.
However, instead of taking the default hyperparameters given in the course, I am trying to perform hyperparameter search using the optuna backend.

Without hp search this code runs perfectly fine for me:
metrics=compute_metrics(start_logits, end_logits, validation_dataset, dataset[“test”])

but once I include it into the hyperparameter search it is throwing an error.
My code looks like this:

and I get the following error for the line where I am extracting the logits.
33 predictions=trainer.predict(validation_dataset)
—> 34 start_logits,end_logits=predictions[0]
35 metrics=compute_metrics(start_logits, end_logits, validation_dataset,dataset[“test”])
36 em=metrics[“exact_match”]
TypeError: cannot unpack non-iterable NoneType object

I think that the trainer.predict() is not returning anything when I am using it in the objective function and I have no ides why since it is working if I use it outside.

Thank you in advance!