| 20 | |
| 21 | |
| 22 | class MemBlock(nn.Module): |
| 23 | def __init__(self, n_in, n_out, act_func): |
| 24 | super().__init__() |
| 25 | self.conv = nn.Sequential(conv(n_in * 2, n_out), act_func, conv(n_out, n_out), act_func, conv(n_out, n_out)) |
| 26 | self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() |
| 27 | self.act = act_func |
| 28 | |
| 29 | def forward(self, x, past): |
| 30 | return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) |
| 31 | |
| 32 | |
| 33 | class TPool(nn.Module): |