How to use 0-3 class of number in MNIST

I would like to know how to use parts 0-3 of the mnist dataset so that the output neurons of the fully connected neural network end up at 4。

here is my code

trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_set = datasets.MNIST(‘…/data’, train=True, transform=trans, download=True)

test_set = datasets.MNIST(‘…/data’, train=False, transform=trans, download=True)

train_loader =, batch_size=batch_size, shuffle=True)
test_loader =, batch_size=batch_size, shuffle=True)

Please help me fix my code

Can you clearly explain your question ? Do you want the outputs of the network to be from 0 to 4 ?

I want to validate 0-3 in the mnist dataset, not 0-9.

You’ll first need to remove the data of other remaining digits from the dataset. After that, simply set the last or the output layer to give four outputs, each corresponding to a digit.

sorry,i dont know how to remove

Take a look at this post which explains you could directly manipulate the internal targets and corresponding data.