(self, channels, use_conv, dims=2)
| 58 | """ |
| 59 | |
| 60 | def __init__(self, channels, use_conv, dims=2): |
| 61 | super().__init__() |
| 62 | self.channels = channels |
| 63 | self.use_conv = use_conv |
| 64 | self.dims = dims |
| 65 | if use_conv: |
| 66 | self.conv = conv_nd(dims, channels, channels, 3, padding=1) |
| 67 | |
| 68 | def forward(self, x): |
| 69 | assert x.shape[1] == self.channels |