I want to replace the encode part of the AAE with a convolutional layer, but it doesn't work anymore. Can you help me? Thanks a lot

Franklin_smith: <p>class Encoder(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(Encoder, self).<strong>init</strong>()</p>
<pre><code>    self.model  = nn.Sequential(
        Conv2dBlock(1,64,(2,2),stride=1,padding=1,norm_fn= 'batchnorm',acti_fn= 'relu'),
        Conv2dBlock(64, 128, (2, 2), stride=1, padding=1, norm_fn='batchnorm', acti_fn='relu'),
         Conv2dBlock(128, 512, (2, 2), stride=1, padding=1, norm_fn='batchnorm', acti_fn='relu')
    )
    self.mu = LinearBlock(512, opt.latent_dim)
    self.logvar = LinearBlock(512, opt.latent_dim)

def forward(self, img):
    x= self.model(img)
    mu = self.mu(x)
    logvar = self.logvar(x)
    z = reparameterization(mu, logvar)
    return z
</code></pre>
<p>class Decoder(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(Decoder, self).<strong>init</strong>()</p>
<pre><code>    self.model = nn.Sequential(
        nn.Linear(opt.latent_dim, 512),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Linear(512, 512),
        nn.BatchNorm1d(512),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Linear(512, int(np.prod(img_shape))),
        nn.Tanh(),
    )

def forward(self, z):
    # print(z.size()) torch.Size([256, 10])
    img_flat = self.model(z)
    # print(img_flat.size()) torch.Size([256, 1024])
    img = img_flat.view(img_flat.shape[0], *img_shape)
    # print(img.size())torch.Size([256, 1, 32, 32])
    return img
</code></pre>
<p>class Discriminator(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(Discriminator, self).<strong>init</strong>()</p>
<pre><code>    self.model = nn.Sequential(
        nn.Linear(opt.latent_dim, 512),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Linear(512, 256),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Linear(256, 1),
        nn.Sigmoid(),
    )

def forward(self, z):
    validity = self.model(z)
    return validity
</code></pre>
<h1>Use binary cross-entropy loss</h1>
<p>adversarial_loss = torch.nn.BCELoss()<br>
pixelwise_loss = torch.nn.L1Loss()</p>
<h1>Initialize generator and discriminator</h1>
<p>encoder = Encoder()<br>
decoder = Decoder()<br>
discriminator = Discriminator()</p>
<p>if cuda:<br>
encoder.cuda()<br>
decoder.cuda()<br>
discriminator.cuda()<br>
adversarial_loss.cuda()<br>
pixelwise_loss.cuda()</p>

ptrblck: <p>What exactly are you changing and what is not working?<br>
Could you post the error message you are seeing?</p>
<p>Also, you can post code snippets by wrapping them into three backticks ```, which makes debugging easier. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

Franklin_smith: <p>Thank you for your reply. I have solved this problem, but I still have a small problem to solve. I have two loss formulas that need to be coded with pytorch, but I have no clue now. Can you help me? Thank you very much<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/5/0/507287c2177bbcc561da26f07346a32cd69a300f.png" alt="微信截图_20200908110737" data-base62-sha1="btFzOC8ZVxYAy89Z2uWz3sMDmlp" width="352" height="62"></p>

Franklin_smith: <p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/2/d21ff1311750262d7a64682cc3e064d219db9314.png" alt="微信截图_20200908110745" data-base62-sha1="tYQHGh96HY4xNTCNdzl8xaMl1Qw" width="285" height="74"></p>