In short :
Although pytorch 1.0 support matrix norm, I still have to use pytorch 0.4.1 due to low cuda version. I implement a nuclear loss function but have a terrible speed maybe because it contains 2 for-loop. How can I make it faster?
In detail :
This is my loss function:
“mat” is my out put which is an 4 dimensional tensor [batch, channel, weight, height]
I want to compute nuclear loss of each matrix in each batch and channel.

nuc.PNG792×255 35.7 KB

With this loss I train my network 10 times slower than before. I want to make it faster but dont know how.
Since torch.mm doesn’t support broadcast, I tried to use torch.matmul but failed.
I would be very grateful if someone teach me how to do.