model.py

import torch
import torch.nn as nn
import torchvision



def create_model(num_classes, seed = 42):
    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
    auto_transforms = weights.transforms()

    model = torchvision.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 = num_classes, bias = True)
    )

    optimizer = torch.optim.Adam(params = model.parameters(), lr = 0.001)
    lossFunc = nn.CrossEntropyLoss()

    return model, auto_transforms, optimizer, lossFunc