fcn.py

class FCN(nn.Module):
    def __init__(self, input_dims, output_dims):
        super(FCN, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dims, 5), 
            nn.LeakyReLU(), 
            nn.Linear(5, output_dims), 
            nn.Sigmoid()
        )
        
    def forward(self, X):
        return self.model(X)