Thank you for the prompt reply!
I realized that error later too. In the end, I was able to transform the classes as wanted, but the code seems really redundant (sorry I’m a newbie!).
Is there a way that I process the conditional transform in the transform() method instead of the __getitem__() method?
Thank you!

ori_transform = transforms.Compose([transforms.Pad(2),
                                transforms.ToTensor(),
                              transforms.Normalize((0.1307,), (0.3081,)),
                              ])
aug_transform = transforms.Compose([
                                transforms.Pad(2),
                                transforms.RandomAffine(degrees=45, translate=(0.1, 0.1), scale=(0.8, 1.2)),
                                transforms.ToTensor(),
                              transforms.Normalize((0.1307,), (0.3081,)),
                              ])



class myMNIST(datasets.MNIST):
    def __init__(self,
                 root,
                 train, 
                 download,
                 transform,):


        super().__init__(root, train, None, None,download)

        self.train= train
        self.transform = transform
        self.download = download
        
        

    def __getitem__(self, index):
        
        x = self.data[index]
        y = self.targets[index]
        aug_classes = [2,3,5,6,7,9]
        if (y.item() in aug_classes):
            

            x = transforms.ToPILImage()(x)           
            x = self.transform(x)

        else:

            x = transforms.ToPILImage()(x)
            x= ori_transform(x)

        return x, y
    
    def transform(self): 
        
        self.transform = aug_transform
        
     

train_dataset_special_aug = myMNIST(root='data', 
                                    train=True,
                                    download=True, 
                                    transform = aug_transform,
                                   )