How to use one class of number in MNIST

You could get the indices for all class1 labels and then index the labels and data:

dataset = datasets.MNIST(root='./data')
idx = dataset.train_labels==1
dataset.train_labels = dataset.train_labels[idx]
dataset.train_data = dataset.train_data[idx]

However, your model won’t learn anything as you just have one class.
Could you explain your use case a bit?
I would at least try to keep two classes in the dataset.

Or do you want to train your model to recognize the number 1 as the valid class and all remaining numbers as false?
If so, you could try this code:

idx = dataset.train_labels != 1
dataset.train_labels[idx] = 0
12 Likes