Solved by fadetoblack in post #4 
              
              
                After this line: 
train_dataset , valid_dataset,  = torch.utils.data.random_split(dataset , [len_train_set, len_valid_set])

You could modify the transforms, something like: 
train_dataset.transforms = train_transform
valid_dataset.transforms = test_transform