Solved by ptrblck in post #2 
              
              
                Try to use a TensorDataset and pass the data in the shape [batch_size, features] to it. 
Once this dataset is created pass it to the DataLoader.