Hi,
According to trace docs

image.png1558×727 34.6 KB

Soooo why don’t you compute the sum of all the elements of the batch and then the trace? Then you would skip the for loop and if i’m not very ****  at this time of the night, it should be equivalent.
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)
        loss = m_total.sum(0).trace()
        loss /= (total_batch*total_channel) #Review this line
        return loss

Regards
Juan