model.py

import torch.nn as nn
import pretrainedmodels as pm

model = pm.__dict__["resnet50"](pretrained='imagenet')

model.avg_pool = nn.AdaptiveAvgPool2d(1)
model.last_linear = nn.Sequential(
    nn.BatchNorm1d(2048),
    nn.Dropout(p=0.25),
    nn.Linear(in_features=2048, out_features=2048),
    nn.ReLU(),
    nn.BatchNorm1d(2048, eps=1e-05, momentum=0.1),
    nn.Dropout(p=0.5),
    nn.Linear(in_features=2048, out_features=1103),
)