model.py

import torch.nn as nn
class SmallNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(8, 64), nn.ReLU())
        self.classifier = nn.Linear(64, 4)
    def forward(self, x):
        return self.classifier(self.encoder(x))