| 63 | |
| 64 | |
| 65 | class PixelNorm(nn.Module): |
| 66 | def __init__(self): |
| 67 | super(PixelNorm, self).__init__() |
| 68 | self.epsilon = 1e-8 |
| 69 | |
| 70 | def forward(self, x): |
| 71 | return x / torch.sqrt(torch.mean(x ** 2, dim=1, keepdim=True) + self.epsilon) |
| 72 | |
| 73 | |
| 74 | class ConvBlock(nn.Module): |