model.py

# creating model.py

########## imports ############
import torch
import torch.nn as nn
from torchvision import models, transforms

###############################

def create_model():
    weights = models.Swin_B_Weights.DEFAULT
    transform = weights.transforms()

    model = models.swin_b(weights = weights)

    for param in model.parameters():
        param.requires_grad = False 

    model.head = nn.Sequential(
        # nn.Linear(in_features = 1024, out_features = 512, bias = True),
        # nn.ReLU(),
        # nn.Dropout(p = 0.3, inplace = True),
        nn.Linear(in_features = 1024, out_features = 101, bias = True)
    )


    return model, transform