Solved by ptrblck in post #4 
              
              
                The workflow to get the predictions sounds correct for the inference use case, but I assume won’t be used to train the model. 
Here is a small example of what I had in mind: 
# setup
batch_size, nb_classes, seq_len = 2, 3, 4
output = torch.randn(batch_size, nb_classes, seq_len, requires_grad=True)
t…