models.py

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as T
from PIL import Image

class LinearClassifier(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearClassifier, self).__init__()

        self.linear_head = nn.Sequential(
            nn.Linear(input_dim, 512, bias=True), 
            nn.ReLU(),
            nn.Linear(512, 256, bias=True),
            nn.ReLU(),
            nn.Linear(256, output_dim, bias=True)
        )

    def forward(self, x):
        return self.linear_head(x)