Fix: extra negatives causes shape mismatch in FactorizedTopK metrics

[patrickorlando]

The Problem

The Retrieval task already slices the candidate_embeddings tensor to remove extra negatives, but it doesn't do the same for the candidate_ids. This leads to a shape mismatch when calculating the FactorizedTopK metrics, if also handling accidental hits.

[caesarjuly]

Thanks for providing this fix. I also noticed this issue. 👍