| 61 | return self.conv(x.reshape(-1, self.stride * C, H, W)) |
| 62 | |
| 63 | class TGrow(nn.Module): |
| 64 | def __init__(self, n_f, stride): |
| 65 | super().__init__() |
| 66 | self.stride = stride |
| 67 | self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False) |
| 68 | def forward(self, x): |
| 69 | _NT, C, H, W = x.shape |
| 70 | x = self.conv(x) |
| 71 | return x.reshape(-1, C, H, W) |
| 72 | |
| 73 | class PixelShuffle3d(nn.Module): |
| 74 | def __init__(self, ff, hh, ww): |