# Different transformations for train and test dataset

**URL:** <https://discuss.pytorch.org/t/different-transformations-for-train-and-test-dataset/158387>\
**Category:** Uncategorized\
**Created:** [August 5, 2022, 4:18am UTC](https://discuss.pytorch.org/t/different-transformations-for-train-and-test-dataset/158387 "2022-08-05T04:18:16Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Ali\_Aghababaei](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ali_aghababaei/32/61342_2.png) [@Ali\_Aghababaei](https://discuss.pytorch.org/u/Ali_Aghababaei)\
**Post date:** [August 5, 2022, 4:18am UTC](https://discuss.pytorch.org/t/different-transformations-for-train-and-test-dataset/158387/1 "2022-08-05T04:18:16Z")

</div>

Hi everyone,  
I have created a custom Dataset class for a number of images, and want to apply K-fold cross-validation on them to split them into train and test datasets (I don’t want to specify a classic validation dataset). The torchvision transforms for train and test datasets are different since I want to apply data augmentation on just the train dataset, not the test dataset. My question is that what is the best way to manage different transformations for each dataset in each fold. should I define a different class for each train or test dataset in each fold?

---

<div class="post-metadata">

**Author:** ![sudomaze](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sudomaze/32/31865_2.png) [@sudomaze](https://discuss.pytorch.org/u/sudomaze)\
**Post date:** [August 6, 2022, 1:11pm UTC](https://discuss.pytorch.org/t/different-transformations-for-train-and-test-dataset/158387/2 "2022-08-06T13:11:01Z")

</div>

Hi Ali,

You can create a function that returns a `Compose` and the parameter passed to it determines if this is a train or test operation. Here is an example:

```python

import torchvision.transforms as T

def get_transforms(mode):
    transforms = []
    if mode == 'train':
        transforms.extend([
            # ...
        ])
    transforms.extend([
        T.Rescale((256,256)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return T.Compose(transforms)

train_set = Dataset(get_transforms('train'))
test_set = Dataset(get_transforms('test'))
# ...

```

You can refer to [this tutorial for more examples](https://pytorch.org/vision/stable/transforms.html).