from torchvision import datasets, transforms

def getData():

    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2023, 0.1994, 0.2010)

    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])

    trainset = datasets.CIFAR10('./data', train=True, download=True, transform=transform_train)
    testset = datasets.CIFAR10('./data', train=False, download=True, transform=transform_test)

    input_dim = 3*32*32
    output_dim = 10
    
    return trainset, testset, input_dim, output_dim
