def __init__()

in models.py [0:0]


    def __init__(self, n_pixels, n_channels, z_dim, distribution="gaussian"):
        super().__init__()

        self.z_dim = z_dim
        self.n_channels = n_channels
        self.distribution = distribution
        self.encoder = nn.Sequential(
            nn.Conv2d(n_channels, n_pixels, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(n_pixels, n_pixels, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(n_pixels, n_pixels, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(n_pixels, 256, kernel_size=1, stride=1),
            nn.ReLU(),
            Lambda(lambda x: x.view(x.size(0), -1)),
            nn.Linear(256, 2 * z_dim),
        )

        self.decoder = nn.Sequential(
            nn.Linear(z_dim, 256),
            nn.ReLU(),
            Lambda(lambda x: x.view(-1, 256, 1, 1)),
            nn.ConvTranspose2d(256, 64, 4),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 64, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, n_pixels, 4, 2, 1),
            nn.ReLU(),
            nn.ConvTranspose2d(n_pixels, n_channels, 4, 2, 1),
            Lambda(lambda x: x.view(x.size(0), -1)),
            nn.ReLU(),
            nn.Linear(32 * 32, n_pixels * n_pixels),
            Lambda(lambda x: x.view(x.size(0), 1, n_pixels, n_pixels)),
            nn.Sigmoid(),
        )