Hi @HiroIshida
Your workaround approach is quite ok.
Why don’t you try:

to inherit WorkaroundDataset from datasets.MNIST,
properly call a constructor of the parent class in the __init__,
then just override __getitem__ to get item from a super's __getitem__

and finally do your stuff with labels before the return of __getitem__