(self, channel, negative_slope=0.2, scale=2 ** 0.5)
| 43 | |
| 44 | class FusedLeakyReLU(nn.Module): |
| 45 | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5): |
| 46 | super().__init__() |
| 47 | self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1)) |
| 48 | self.negative_slope = negative_slope |
| 49 | self.scale = scale |
| 50 | |
| 51 | def forward(self, input): |
| 52 | return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale) |