Solved by vdw in post #2 
              
              
                What do you mean the batch is no longer represented? It’s in the second dimension :). If you need it in the first one you could simply do output.transpose(0,1)