It seems that the return value of prediction_step needs to be returned as numpy.
def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
with torch.no_grad():
loss, out = self.compute_loss(model, inputs, return_outputs=True)
a, p, n = out.embeddings
preds = tuple(t.detach().cpu().numpy() for t in (a, p, n)) # numpy for metrics
return (loss.detach(), preds, None) # label_ids can be None