How to use one class of number in MNIST

I think:

dataset.train_labels = dataset.train_labels[idx]

is giving “AttributeError: can’t set attribute.” Error.
This will be:
train_labels = dataset.train_labels[idx]

Anyway, thank you for sharing the same.