net.py

class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()
        self.NeurlNet = nn.Sequential(
            nn.Linear(50, 40),
            nn.ReLU(),
            nn.Linear(40, 30),
            nn.ReLU(),
            nn.Linear(30, 1)
        )
    
    def forward(self, input_, apply_sigmoid=False):
        x = self.NeurlNet(input_)
        x = x.squeeze()
        if apply_sigmoid:
            x = torch.sigmoid(x)
        
        return x