| 3 | |
| 4 | |
| 5 | class ScalingLayer(nn.Module): |
| 6 | def __init__(self, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]): |
| 7 | super().__init__() |
| 8 | self.register_buffer('shift', torch.Tensor(mean)[None, :, None, None]) |
| 9 | self.register_buffer('scale', torch.Tensor(std)[None, :, None, None]) |
| 10 | |
| 11 | def forward(self, inp): |
| 12 | return (inp - self.shift) / self.scale |
| 13 | |
| 14 | def inv(self, inp): |
| 15 | return inp * self.scale + self.shift |
no outgoing calls
no test coverage detected