Untitled

import torch.nn as nn
import torchvision.models as models

class ResNet50WithDropout(nn.Module):
    def __init__(self, num_classes=len(class_names), dropout=0.5):
        super(ResNet50WithDropout, self).__init__()
        self.resnet50 = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        num_ftrs = self.resnet50.fc.in_features
        self.resnet50.fc = nn.Identity()
        self.fc = nn.Sequential(
            nn.Dropout(p=dropout),
            nn.Linear(num_ftrs, len(class_names)),
            nn.CELU()
        )

    def forward(self, x):
        x = self.resnet50(x)
        x = self.fc(x)
        return x