import torch import torchvision from torch import nn device='cpu' def create_model(num_classes: int=13): weights = torchvision.models.EfficientNet_B0_Weights.DEFAULT transforms = weights.transforms() model = torchvision.models.efficientnet_b0(weights=weights).to(device) for param in model.features.parameters(): param.requires_grad = False model.classifier = torch.nn.Sequential( torch.nn.Dropout(p=0.2, inplace=True), torch.nn.Linear(in_features=1280, out_features=13, bias=True) ).to(device) return model, transforms