a question for JMMD loss computation tllib/alignment/dan.py#113-114

[qwerdfjklol]

When using the Joint Maximum Mean Discrepancy (JMMD) loss, I observed that the loss value became negative during the later stages of training. Upon further investigation, I discovered that this issue arose because the calculations did not account for the diagonal elements. Could you please explain the reason why you ignore the diagonal element during computation?

index_matrix = torch.zeros(2 * batch_size, 2 * batch_size)

for i in range(batch_size):
    for j in range(batch_size):
        if i != j:
            index_matrix[i][j] = 1. / float(batch_size * (batch_size - 1))
            index_matrix[i + batch_size][j + batch_size] = 1. / float(batch_size * (batch_size - 1))
for i in range(batch_size):
    for j in range(batch_size):
        index_matrix[i][j + batch_size] = -1. / float(batch_size * batch_size)
        index_matrix[i + batch_size][j] = -1. / float(batch_size * batch_size)