Fix logits compuation in KTO trainer prediction step

[issamemari]

Description of the issue

There is a bug in the following few lines of code in kto_trainer.py

logits_dict = {
    "eval_logits/chosen": metrics["logits/chosen"],
    "eval_logits/rejected": metrics["logits/rejected"],
}
logits = tuple(v.unsqueeze(dim=0) for k, v in logits_dict.items() if k not in ignore_keys)
logits = torch.stack(logits).mean(axis=1).to(self.accelerator.device)

This assumes that the values in the logits_dict are tensors, but they are not. These values are computed in get_batch_loss_metrics where the the logits are averaged and .item() is called on the resulting tensor to get a float.

This causes the following error when running a KTO training:

AttributeError: 'float' object has no attribute 'unsqueeze'

What does this PR do?

• Treat the values of logits_dict as floats
• Fixes AttributeError: 'float' object has no attribute 'unsqueeze' during the evaluation phase of a KTO training

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.