| 44 | |
| 45 | |
| 46 | class Block(nn.Module): |
| 47 | def __init__(self, in_channels: int, out_channels: int, stride: int, depth: int, dropout: float): |
| 48 | super(Block, self).__init__() |
| 49 | self.block = nn.Sequential( |
| 50 | DownsampleUnit(in_channels, out_channels, stride, dropout), |
| 51 | *(BasicUnit(out_channels, dropout) for _ in range(depth)) |
| 52 | ) |
| 53 | |
| 54 | def forward(self, x): |
| 55 | return self.block(x) |
| 56 | |
| 57 | |
| 58 | class WideResNet(nn.Module): |