(self, x: torch.Tensor)
| 154 | ) |
| 155 | |
| 156 | def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: |
| 157 | x = self.convInit(x) |
| 158 | if self.dropout_prob is not None: |
| 159 | x = self.dropout(x) |
| 160 | |
| 161 | down_x = [] |
| 162 | |
| 163 | for down in self.down_layers: |
| 164 | x = down(x) |
| 165 | down_x.append(x) |
| 166 | |
| 167 | return x, down_x |
| 168 | |
| 169 | def decode(self, x: torch.Tensor, down_x: list[torch.Tensor]) -> torch.Tensor: |
| 170 | for i, (up, upl) in enumerate(zip(self.up_samples, self.up_layers)): |