Upload modeling_vae.py

modeling_vae.py

@@ -29,7 +29,7 @@ class VAEModel(PreTrainedModel):
             outputDim=self.hidden_channels * (2 ** i)
             self.encoder.append(nn.Conv2d(inputDim, outputDim, kernel_size=4, stride=2, padding=1))# -> (hidden, H/2, W/2)
             self.encoder.append(nn.BatchNorm2d(outputDim))
-            self.encoder.append(self.actFn())
+            self.encoder.append(nn.ReLU())
         self.encoderD = nn.Sequential(*self.encoder)
 
         with torch.no_grad():
@@ -49,7 +49,7 @@ class VAEModel(PreTrainedModel):
             outputDim=self.hidden_channels * i
             self.decoder.append(nn.ConvTranspose2d(inputDim, outputDim, kernel_size=4, stride=2, padding=1))
             self.decoder.append(nn.BatchNorm2d(outputDim))
-            self.decoder.append(self.actFn())
+            self.decoder.append(nn.ReLU())
             inputDim=outputDim
         H_before_last = self.enc_H * (2 ** (self.encoder_layers - 1))
         W_before_last = self.enc_W * (2 ** (self.encoder_layers - 1))