from pathlib import Path

import torch
from torch import nn
from torch.export import Dim


class ScoreModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(4, 8),
            nn.ReLU(),
            nn.Linear(8, 2),
        )

    def forward(self, features):
        return torch.softmax(self.layers(features), dim=-1)