model.py

import torchvision
from torch import nn

def create_vit_b_16_swag(num_classes:int=101, seed:int=42):
    # 1. Get the base mdoel with pretrained weights and send to target device
    weights = torchvision.models.ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1
    transforms = weights.transforms()
    model = torchvision.models.vit_b_16(weights=weights)#.to(device)

    # 2. Freeze the base model layers
    for param in model.parameters():
        param.requires_grad = False


    # 3. Change the heads
    model.heads = nn.Linear(in_features=768,
                            out_features=101)#.to(device)

    # 5. Give the model a name
    model.name = "vit_b_16_swag"
    print(f"[INFO] Created new {model.name} model.")
    return model, transforms