def __init__(self, input_dim, cat_dim, co_cat_dim):
        super(Discriminator, self).__init__()

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            # if normalize:
            #     layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.dis_model = nn.Sequential(
            *block(input_dim, 32, normalize=False),
            *block(32, 32),
        )

        # Output layers
        self.discriminator = nn.Sequential(
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

        self.encoder = nn.Sequential(
            nn.Linear(32, cat_dim),
            nn.Softmax()
        )

        self.co_encoder = nn.Sequential(
            nn.Linear(32, co_cat_dim),
            nn.Softmax()
        )

    def forward(self, input):
        out = self.dis_model(input)
        validity = self.discriminator(out)
        label = self.encoder(out)
        colabel = self.co_encoder(out)

        return validity, label, colabel

and loss function is just
def mse_loss(output, target):
    return nn.MSELoss()(output, target)