diffusions5M.py

import timm
import torch
from torch import nn

class Model5M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = timm.create_model('timm/tf_mobilenetv3_large_100.in1k', pretrained=False, num_classes=0)

        self.clf = nn.Sequential(
            nn.Linear(1280, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 2))

    def forward(self, image):
        image_features = self.model(image)
        return self.clf(image_features)