Solved by ptrblck in post #4 
              
              
                This would work: 
output = ... # shape [batch_size, seq_len, nb_classes]
output = output.permute(0, 2, 1).contiguous() # shape [batch_size, nb_classes, seq_len]