| 43 | return x |
| 44 | |
| 45 | class Decoder(nn.Module): |
| 46 | def __init__(self) -> None: |
| 47 | super().__init__() |
| 48 | self.fc = nn.Linear(LATENT_DIMS, 16*60*60) |
| 49 | self.conv2 = nn.ConvTranspose2d(16, 6, 3) |
| 50 | self.conv1 = nn.ConvTranspose2d(6, 1, 3) |
| 51 | self.relu = nn.ReLU() |
| 52 | |
| 53 | def forward(self, x): |
| 54 | x = self.fc(x) |
| 55 | x = x.view(-1, 16, 60, 60) # infer first dim from other dims |
| 56 | x = self.conv2(x) |
| 57 | x = self.relu(x) |
| 58 | x = self.conv1(x) |
| 59 | x = self.relu(x) |
| 60 | return x |
| 61 | |
| 62 | class Autoencoder(nn.Module): |
| 63 | def __init__(self) -> None: |