import torchvision

import dataset


class Cifar10Augmented(dataset.Cifar10):

    NAME = "cifar10-augmented"

    # === PROTECTED ===

    def create_training_augmentation(self):
        return [
            torchvision.transforms.RandomCrop(32, padding=4),
            torchvision.transforms.RandomHorizontalFlip()
        ]
