| 39 | return torch.tanh(x / 3) * 3 |
| 40 | |
| 41 | class MemBlock(nn.Module): |
| 42 | def __init__(self, n_in, n_out): |
| 43 | super().__init__() |
| 44 | self.conv = nn.Sequential( |
| 45 | conv(n_in * 2, n_out), nn.ReLU(inplace=True), |
| 46 | conv(n_out, n_out), nn.ReLU(inplace=True), |
| 47 | conv(n_out, n_out) |
| 48 | ) |
| 49 | self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() |
| 50 | self.act = nn.ReLU(inplace=True) |
| 51 | def forward(self, x, past): |
| 52 | return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) |
| 53 | |
| 54 | class TPool(nn.Module): |
| 55 | def __init__(self, n_f, stride): |