What to do with with a tensor(1d array of) cosine similarity values

I am performing cosine similarity (nn.cosineSimilarity()) between two 2D tensors (of same shape of course). Now, the resultant output is a 1D tensor which contains n single tensors. These single tensors are the pairwise cosine similarities.

Now, my question what can I do with these pairwise cosine similarities. For training, I am passing them directly to my custom loss function and things seems to work fine.

For Prediction

I have M such outputs out of which I have to choose the one with max score/output. But how can I select a max from arrays.

My approach for :
I am flattening the 2d tensors (which are outputs from a GRU) and then calculating the cosine sim thus, getting a single tensor. This approach seems to work but the accuracy of the model is not as expected.
Hence, I am looking for other alternatives.

Another thing I had tried was to just select the array with the maximum single tensor among all.

What can I do with an array of pairwise similarities?