| 269 | |
| 270 | |
| 271 | class NoiseInjection(nn.Module): |
| 272 | def __init__(self): |
| 273 | super().__init__() |
| 274 | |
| 275 | self.weight = nn.Parameter(torch.zeros(1)) |
| 276 | |
| 277 | def forward(self, image, noise=None): |
| 278 | |
| 279 | if noise is None: |
| 280 | return image |
| 281 | else: |
| 282 | return image + self.weight * noise |
| 283 | |
| 284 | |
| 285 | class ConstantInput(nn.Module): |