| 31 | |
| 32 | |
| 33 | class TPool(nn.Module): |
| 34 | def __init__(self, n_f, stride): |
| 35 | super().__init__() |
| 36 | self.stride = stride |
| 37 | self.conv = nn.Conv2d(n_f * stride, n_f, 1, bias=False) |
| 38 | |
| 39 | def forward(self, x): |
| 40 | _NT, C, H, W = x.shape |
| 41 | return self.conv(x.reshape(-1, self.stride * C, H, W)) |
| 42 | |
| 43 | |
| 44 | class TGrow(nn.Module): |