Say I have C classes, and I want to get a C*C similarity matrix in which each entry is the cosine similarity between class_i and class_j. I write the below code to compute the similar loss based on the weights of last but one fc layer. Below is the part of the code for simplicity：
cos = nn.CosineSimilarity(dim=1, eps=1e-6)
for batch_idx, (data, target) in enumerate(self.data_loader):
    # C*M
    weights = self.model.module.model.fc[-1].weight
    num_iter = weights.size(0)
    # similarity matrix
    sim_mat = torch.empty(0, num_iter).cuda()
    for j in range(num_iter):
           weights_i = weights[j, :].expand_as(weights)
           sim = cos(weights, weights_i)
           sim = torch.unsqueeze(sim, 0)
           sim_mat = torch.cat((sim_mat, sim), 0)
    sim_mat = sim_mat - torch.diag(torch.diag(sim_mat))

    max_val = torch.max(sim_mat, dim=1).values

    loss_sim = (torch.sum(max_val) / num_classes**2)
    ...
    loss_sim.backward()  

Why my loss_sim doesn’t work for my training process? It seems that the loss_sim didn’t backward properly to affect the original model weights.
I know we can wrap the loss function into a class like:
class Loss_sim(nn.Module):
    def __init__():

    def forward(self, weights):

Should I write like this? Is any problem here because of the copy of the original weights?
Thanks in advance