If you pass my model the batch ([0,1,2],[0,1,2]), then it will output a scalar for each of the case 0-0, 1-1 and 2-2, i.e. my output will be 3 scalars.



 ptrblck:

In the former case, you could increase the inner batch ( j ) by one while increasing the outer batch by batch_size .


As I have mentioned in my question, I am using batch size of 8. I replicate the i 8 times (for each element) and pass it to 8 elements of the inner batch, j. But, this is too slow. I want to optimise this part.



 ptrblck:

You could add a check for the indices as:
for index_i, i in enumerate(A):
    temp = []
    for index_j, j in enumerate(B):
        if j > i: # or j >= i, if you don't want to compare the same indices
            continue



Yes, I could do that, but, is there a way to combine this with the previous optimisation (if found) to further speed up the task.