(self, x, skip=None)
| 164 | self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, 2, stride=2) |
| 165 | |
| 166 | def forward(self, x, skip=None): |
| 167 | if self.down: |
| 168 | x = self.conv(x) |
| 169 | pool = self.pool(x) |
| 170 | return x, pool |
| 171 | else: |
| 172 | x = self.up(x) |
| 173 | if skip is not None: |
| 174 | x = torch.cat([x, skip], dim=1) |
| 175 | x = self.conv(x) |
| 176 | return x |
| 177 | |
| 178 | class ConvNeXtBlock(nn.Module): |
| 179 | def __init__(self, dim, drop_path=0.0, layer_scale_init_value=1e-6): |