Solved by ptrblck in post #3 
              
              
                in = torch.cat((in, ...)) will slow down your code as you are concatenating to the same tensor in each iteration. Append to data to a list and create the tensor after all samples of the current batch were already appended to it.