import torch from torchvision import models from torch import nn def create_model(num_classes:int=3): Eff_Net_Model=models.efficientnet_b2(weights="DEFAULT") Eff_Net_Weights=models.EfficientNet_B2_Weights.DEFAULT Eff_Net_transform=Eff_Net_Weights.transforms() for param in Eff_Net_Model.features.parameters(): param.requires_grad=False Eff_Net_Model.classifier=nn.Sequential( nn.Dropout(p=0.3,inplace=True), nn.Linear(in_features=1408, out_features=num_classes, bias=True) ) return Eff_Net_Model,Eff_Net_transform