[FEA] Optimize `torch_gather_like` operator used in top-k extraction

[sararb]

🚀 Feature request

Motivation

• Create a vectorized version of torch_gather_like function implemented here. This function is used to gather top-k indices from each row of one-hot vector of labels.

Your contribution

Maybe later check if we could come up with a more efficient implementation of gather_torch_like which does not using TensorArray, but rather uses only tf.gather_nd() with a 2D index (rows, cols), like the one I implemented here here.
Not a priority for now, but maybe it might improve the metrics performance

Originally posted by @gabrielspmoreira in #189 (comment)

[sararb]

For context, this is the expected output of all_indices defined in torch_gather_like function where for each example in the batch (size of 16 in this example), we are building the 2-d indices that allow us to retrieve the one-hot value related to the top-k (k=5 in this example) positions.