Can I custom transform MNIST images and targets in transforms.Compose?

I am trying to flatten MNIST images and convert labels to one hot encoded vectors in transforms.Compose itself. How can that be done?

The MNIST dataset from torchvision is in PIL image. So if you want to flatten MNIST images, you should transform the images into tensor format by transforms.ToTensor() in transforms.Compose([]). And then you could use DataLoader to load the images, read and flatten batches of them.

In my shallow view, you can not convert the sparse labels into one-hot format in transforms.Compose. You could change them into one-hot label per batch by this:

one_hot = torch.zeros(len(sparse_label), 10)
one_hot = one_hot.scatter_(1, sparse_label.unsqueeze(1), 1)

@MariosOreo Thanks for your answer.

Yeah I saw that there is no way to get “target labels” inside transforms.Compose since the call to it includes “img” only. I was accessing MNIST dataset’s train_data and train_labels. I thought transforms would apply to them as well. My bad. I later found that they are applied only in data loaders.