Whats the proper way to push all data to GPU and then take small batches during training?

For MNIST, the dataset already stores everything in Tensors, so you can grab the ds.data and ds.targets Tensors from the torchvision MNIST dataset directly and stick them into your TensorDataset.

Best regards

Thomas

1 Like