Correct data loading, splitting and augmentation in Pytorch
The tutorial doesn't seem to explain how we should load, split and do proper augmentation.

Let's have a dataset consisting of cars and cats. The folder structure would be:

data
  cat
    0101.jpg
    0201.jpg
    ...
  dogs
    0101.jpg
    0201.jpg
    ...


At first, I loaded the dataset by datasets.ImageFolder function. Image Function has command "TRANSFORM" where we can set some augmentation commands, but we don't want to apply augmentation to test dataset! So let's stay with transform=None. 

data = datasets.ImageFolder(root='data')


Apparently, we don't have folder structure train and test and therefore I assume a good approach would be to use split_dataset function 

    train_size = int(split * len(data))
    test_size = len(data) - train_size
    train_dataset, test_dataset = torch.utils.data.random_split(data, [train_size, test_size])


Now let's load the data the following way. 

    train_loader = torch.utils.data.DataLoader(train_dataset,
                                              batch_size=8,
                                              shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=8,
                                              shuffle=True)


How can I apply transformations (data augmentation) to the "train_loader" images? 

Basically I need to: 1. load data from the folder structure explained above
2. split the data into test/train parts
3. apply augmentations on train part.  

I am not sure if there is a recommended way of doing this, but this is how I would workaround this problem:

Given that torch.utils.data.random_split() returns Subset, we cannot (can we? not 100% sure here I double-checked, we cannot) exploit their inner datasets, because they are the same (the only diference is in the indices). In this context, I would implement a simple class to apply transformations, something like this:

from torch.utils.data import Dataset

class ApplyTransform(Dataset):
    """
    Apply transformations to a Dataset

    Arguments:
        dataset (Dataset): A Dataset that returns (sample, target)
        transform (callable, optional): A function/transform to be applied on the sample
        target_transform (callable, optional): A function/transform to be applied on the target

    """
    def __init__(self, dataset, transform=None, target_transform=None):
        self.dataset = dataset
        self.transform = transform
        self.target_transform = target_transform
        # yes, you don't need these 2 lines below :(
        if transform is None and target_transform is None:
            print("Am I a joke to you? :)")

    def __getitem__(self, idx):
        sample, target = self.dataset[idx]
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return sample, target

    def __len__(self):
        return len(self.dataset)


And then use it before passing the dataset to the dataloader:

import torchvision.transforms as transforms

train_transform = transforms.Compose([
    transforms.ToTensor(),
    # ...
])
train_dataset = ApplyTransform(train_dataset, transform=train_transform)

# continue with DataLoaders...

