from torch import nn
import torch

class OurModel(nn.Module):
    def __init__(self, num_labels, emb_dim=10):
        super(OurModel, self).__init__()
        self.emb = nn.Embedding(num_labels, emb_dim)
        self.user_embs = nn.Embedding(1, emb_dim)

    def forward(self, pos, neg):
        batch_size = neg.size(0)
        one_user_vector = self.user_embs(torch.zeros(1).long())
        repeated_user_vector = one_user_vector.repeat((batch_size, 1)).view(batch_size, -1, 1)
        pos_res = torch.bmm(self.emb(pos), repeated_user_vector).squeeze(2)
        neg_res = torch.bmm(self.emb(neg), repeated_user_vector).squeeze(2)

        return pos_res, neg_res
        
num_labels = 100
model = OurModel(num_labels)

pos_labels = torch.randint(high=num_labels, size=(3,1)) # our five labels
neg_labels = torch.randint(high=num_labels, size=(3,2)) # a few random negatives per positive

pos_res, neg_res = model(pos_labels, neg_labels)
print('Positive Labels:', pos_labels)
print('Negative Labels:', neg_labels)
print('Model positive scores:', pos_res)
print('Model negative scores:', neg_res)
loss = warp_loss(pos_res, neg_res, num_labels=num_labels, device=torch.device('cpu'))
print('Loss:', loss)
loss.backward()

Positive Labels: tensor([[21],
        [71],
        [26]])
Negative Labels: tensor([[47, 10],
        [56, 78],
        [44, 55]])
Model positive scores: tensor([[-4.9562],
        [-1.6886],
        [ 3.3984]], grad_fn=<SqueezeBackward1>)
Model negative scores: tensor([[ 1.0491,  4.9357],
        [-2.1289,  0.4496],
        [ 3.4541,  0.0931]], grad_fn=<SqueezeBackward1>)
Loss: tensor(39.6134, grad_fn=<SumBackward0>)

print('We can also see that the gradient is only active for 2x the number of positive labels:', (model.emb.weight.grad.sum(1) != 0).sum().item())
print('Meaning we correctly discard the gradients for all other than the offending negative label.')

We can also see that the gradient is only active for 2x the number of positive labels: 6
Meaning we correctly discard the gradients for all other than the offending negative label.