"""
How to load a network using safetensors:

```python
from safetensors.torch import load_model

class MysteryNet(nn.Module):
    ...

model = MysteryNet()
 _, _ = load_model(model=model, filename="model.safetensors")

```
"""

import torch.nn as nn
import torch.nn.functional as F


class MysteryNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.dense1 = nn.Linear(in_features=2, out_features=16)
        self.dense2 = nn.Linear(in_features=16, out_features=32)
        self.bn2 = nn.BatchNorm1d(32)

        self.dense3 = nn.Linear(in_features=32, out_features=32)
        self.dense4 = nn.Linear(in_features=32, out_features=32)
        self.bn4 = nn.BatchNorm1d(32)

        self.dense5 = nn.Linear(in_features=32, out_features=32)
        self.dense6 = nn.Linear(in_features=32, out_features=32)
        self.bn6 = nn.BatchNorm1d(32)

        self.dense7 = nn.Linear(in_features=32, out_features=32)
        self.dense8 = nn.Linear(in_features=32, out_features=16)
        self.bn8 = nn.BatchNorm1d(16)

        self.dense_head = nn.Linear(in_features=16, out_features=3)

    def forward(self, x):
        x = self.dense1(x)
        x = F.gelu(x)
        x = self.dense2(x)
        x = F.gelu(x)
        x = self.bn2(x)

        x = self.dense3(x)
        x = F.gelu(x)
        x = self.dense4(x)
        x = F.gelu(x)
        x = self.bn4(x)

        x = self.dense5(x)
        x = F.gelu(x)
        x = self.dense6(x)
        x = F.gelu(x)
        x = self.bn6(x)

        x = self.dense7(x)
        x = F.gelu(x)
        x = self.dense8(x)
        x = F.gelu(x)
        x = self.bn8(x)

        x = self.dense_head(x)

        return x