Use a for loop in loss calculation is a little bit slow.

[sagittahjz]

Use a for loop in loss calculation is a little bit slow.
You can find a way to remove the for loop.
In my case, only pairs on the diagonal are positive, so I remove the for loop as follows.

simi_mat = torch.matmul(y1, torch.t(y2))
simi_sub = simi_mat - ms_gama
pos_pair_sub = torch.unsqueeze(torch.diag(simi_sub), 1)
neg_pair_sub_plus1 = simi_sub
neg_pair_sub_plus1[range(batch_size), range(batch_size)] = 0
pos_loss = torch.log(1 + torch.sum(torch.exp(-ms_alpha * pos_pair_sub), dim = 1)) / ms_alpha
neg_loss = torch.log(torch.sum(torch.exp(ms_beta * neg_pair_sub_plus1), dim = 1)) / ms_beta
loss = torch.mean(pos_loss + neg_loss)