transforms.py

from torchvision import transforms, models
from typing import Literal, Dict

_weights = models.ViT_B_16_Weights.DEFAULT

model_transforms: Dict[Literal['Custom', 'Pretrained'], transforms.Compose] = {
    'Custom': transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ]),
    'Pretrained': _weights.transforms()
}