(self, z)
| 131 | self.bn2 = nn.BatchNorm(num_filters_3) |
| 132 | |
| 133 | def __call__(self, z): |
| 134 | x = self.lin1(z) |
| 135 | |
| 136 | # reshape to BHWC |
| 137 | x = x.reshape( |
| 138 | -1, self.input_shape[0], self.input_shape[1], self.max_num_filters |
| 139 | ) |
| 140 | |
| 141 | # approximate transposed convolutions with nearest neighbor upsampling |
| 142 | x = nn.leaky_relu(self.bn1(self.upconv1(x))) |
| 143 | x = nn.leaky_relu(self.bn2(self.upconv2(x))) |
| 144 | # sigmoid to ensure pixel values are in [0,1] |
| 145 | x = mx.sigmoid(self.upconv3(x)) |
| 146 | return x |
| 147 | |
| 148 | |
| 149 | class CVAE(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected