Thanks for the reply.
I do have an idea of how I want to implement this, but it is limited to just one batch. If I flatten the all batches and then use that function, the search space for comparison will be too large; since each duplicate value will be compared with all the other values in flattened tensor. Is there away I can apply a function to all batches individually?