cifar_datamodule.py

cifar_datamodule.py

from flash.image import ImageClassifier, ImageClassificationData
class CIFAR10DataModule(ImageClassificationData):
    
    @property
    def num_classes(self):
        return 10
      
dm = CIFAR10DataModule.from_datasets(
    train_dataset=train_set,
    test_dataset=test_set,
    train_transform=train_transforms,
    test_transform=test_transforms,
    # Do not forget to set `predict_transform`,
    # this is what we will use for uncertainty estimation!
    predict_transform=test_transforms,
    batch_size=64,
)