text: import torch.nn as nn
  from utils.model import weights_init
  
  
  class Discriminator(nn.Module):
      """
      Input shape: (N, 3, 64, 64)
      Output shape: (N, )
      """
      def __init__(self, in_dim, dim=64):
          super(Discriminator, self).__init__()
  
          def conv_bn_lrelu(in_dim, out_dim):
              return nn.Sequential(
                  nn.Conv2d(in_dim, out_dim, 5, 2, 2),
                  nn.BatchNorm2d(out_dim),
                  nn.LeakyReLU(0.2),
              )
  
          """ Medium: Remove the last sigmoid layer for WGAN. """
          self.ls = nn.Sequential(
              nn.Conv2d(in_dim, dim, 5, 2, 2),
              nn.LeakyReLU(0.2),
              conv_bn_lrelu(dim, dim * 2),
              conv_bn_lrelu(dim * 2, dim * 4),
              conv_bn_lrelu(dim * 4, dim * 8),
              nn.Conv2d(dim * 8, 1, 4),
              # nn.Sigmoid(),
          )
          self.apply(weights_init)
  
      def forward(self, x):
          y = self.ls(x)
          y = y.view(-1) # y.shape [64, 1, 1, 1]
          return y
  