Question on "in batch softmax" and retrieval task

Hi,

Thanks all for this awesome library ! I have a question regarding the retrieval task and the way it's implemented. From what I understand from the source code, the retrieval task uses CategoricalCrossEntropy under the hood with a queries x candidates embeddings matrix multiplication to get the scores. All of these are compared with an identity matrix which acts as the y_true variable in the loss function.

Let's assume, as a very simple example, that we have 2 users and 2 products, and that the batch size is 2. Let's now assume our dataset looks like :

For the first batch, the network will work to get the dot product of the embeddings of A and 1 close to 1, and the dot product of A and 2 close to 0 (cf identity matrix).
For the second batch, the network will work to get the dot product of the embeddings of A and 2 close to 1, and the embeddings of A and 1 close to 0, which is the exact opposite of what the network was told at batch 1.

My question is therefore : how can we avoid "erasing" or "contradicting" interactions for users who come back in different batches ?

Thanks a lot for your help !