ptbp1_model.py

class MoonsModel(nn.Module):
    def __init__(self, n_features, n_neurons):
        super(MoonsModel, self).__init__()
        self.hidden = nn.Linear(in_features=n_features, out_features=n_neurons)
        self.out_layer = nn.Linear(in_features=n_neurons, out_features=2)

    def forward(self, X):
        out = F.relu(self.hidden(X))
        out = F.sigmoid(self.out_layer(out))
        return out