This is the modified version. Although it is getting faster, it is still very time consuming. But it can only be like this, I guess. Thank you very much for helping me solve this problem.
class NuclearLossFunc(nn.Module):
    def __init__(self):
        super(NuclearLossFunc, self).__init__() 
        return
    def forward(self,mat):
        loss = torch.zeros([1]).cuda()
        total_batch, total_channel, Sx, Sy = mat.size()
        mat = mat.view(-1,Sx,Sy) 
        mat_trans = torch.transpose(mat,1,2)
        m_total = torch.bmm(mat_trans,mat)
        for i in range(m_total.size()[0]): 
            loss += m_total[i,:,:].trace()
        loss /= m_total.size()[0] 
        return loss

p.s. I didn’t use sqrt because it’s too slow. You will need a smaller weight to this loss function.