pytorch: sum of cross entropy over all classes
I want to compute sum of cross entropy over all classes for each prediction, where the input is batch (size n), and the output is batch (size n). 



The simplest way is for loop (for 1000 classes):

def sum_of_CE_lost(input):
    L = 0
    for c in range(1000):
        L = L + torch.nn.CrossEntropyLoss(input, c)
    return L


However, it is very slow. What is a better way? How can we parallelized it for GPU (CUDA)?

I found the answer:

torch.nn.functional.log_softmax (input).sum() / input.shape[0]


We divide by input.shape[0] because cross_entropy() takes, by default the mean across the batch dimension.
