(self, x)
| 36 | self.opt = opt |
| 37 | |
| 38 | def forward(self, x): |
| 39 | if x.size(2) != 256 or x.size(3) != 256: |
| 40 | x = F.interpolate(x, size=(256, 256), mode='bilinear') |
| 41 | |
| 42 | x = self.layer1(x) |
| 43 | x = self.layer2(self.actvn(x)) |
| 44 | x = self.layer3(self.actvn(x)) |
| 45 | x = self.layer4(self.actvn(x)) |
| 46 | x = self.layer5(self.actvn(x)) |
| 47 | if self.opt.crop_size >= 256: |
| 48 | x = self.layer6(self.actvn(x)) |
| 49 | x = self.actvn(x) |
| 50 | |
| 51 | x = x.view(x.size(0), -1) |
| 52 | mu = self.fc_mu(x) |
| 53 | logvar = self.fc_var(x) |
| 54 | |
| 55 | return mu, logvar |