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.efficientnet_b2(weights = weights)

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

    model.classifier = nn.Sequential(
        nn.Dropout(p = 0.3, inplace = True),
        nn.Linear(in_features = 1408, out_features = 101, bias = True)
    )


    return model, transform