Also some utils:
def create_transforms(input_size):
    channel_means = [0.485, 0.456, 0.406]
    channel_stds = [0.229, 0.224, 0.225]

    train_tfms = transforms.Compose([transforms.Resize(input_size),
                                     transforms.ToTensor(),
                                     transforms.Normalize(channel_means, channel_stds)])
    test_tfms = transforms.Compose([transforms.Resize(input_size),
                                    transforms.ToTensor(),
                                    transforms.Normalize(channel_means, channel_stds)])
    mask_tfms = transforms.Compose([transforms.Resize(input_size),
                                    transforms.ToTensor()])
    return train_tfms, test_tfms, mask_tfms

def create_dataloaders(data_dir, input_size=256, test_pct=0.2, batch_size=64) -> (DataLoader, DataLoader, DataLoader):
    train_data, test_data, eval_data = load_datasets(data_dir, input_size, test_pct)

    trainloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=6, pin_memory=True)
    testloader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=6, pin_memory=True)
    evalloader = DataLoader(eval_data, batch_size=1, shuffle=False, num_workers=6, pin_memory=True)
    return trainloader, testloader, evalloader

trainloader, testloader, evalloader = create_dataloaders(data_dir, test_pct=test_pct, batch_size=batch_size)