def wasserstein_loss(real_logit, fake_logit):
    """
    :param real_logit: logit(s) for real images (if None just return generator loss)
    :param fake_logit: logit(s) for fake images
    :return: loss for discriminator and generator (unless real_logit is None)
    """
    loss_generator = - fake_logit
    if real_logit is None:
        return loss_generator
    loss_discriminator_real = - real_logit
    loss_discriminator_fake = fake_logit
    #  this actually negates the need for a bias in the FC layer, it's cancelled out
    loss_discriminator = loss_discriminator_real + loss_discriminator_fake
    return loss_discriminator, loss_generator