import torchvision import torch def create_model( num_classes: int=4, seed: int=42): weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT transform = weights.transforms() model = torchvision.models.efficientnet_b2(weights=weights).to("cpu") for params in model.parameters(): params.requires_grad = False model.classifier = torch.nn.Sequential( torch.nn.Dropout(p=0.2,inplace=True), torch.nn.Linear(1408,num_classes) ) return model, transform