import numpy as np
import onnx
import onnxruntime as ort
import torch


class SmallClassifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(4, 8),
            torch.nn.ReLU(),
            torch.nn.Linear(8, 3),
        )

    def forward(self, features):
        return self.layers(features)


torch.manual_seed(7)
model = SmallClassifier().eval()
example_input = torch.randn(1, 4)