Solved by richard in post #3 
              
              
                It’s standard for pytorch tensors to be organized in a (N, C, <other_sizes>) fashion. For example, images would be represented as (N, C, H, W). 
You could permute the dimensions to use NLLLoss: 
tensor # Nxd1xd2…dkxC
tensor.permute(0, 2, 3, ..., 1)