| 58 | } |
| 59 | |
| 60 | torch::nn::Sequential CenterBlock(int in_channels, int out_channels){ |
| 61 | return torch::nn::Sequential(Conv2dReLU(in_channels, out_channels, 3, 1), |
| 62 | Conv2dReLU(out_channels, out_channels, 3, 1)); |
| 63 | } |
| 64 | |
| 65 | UNetDecoderImpl::UNetDecoderImpl(std::vector<int> encoder_channels, std::vector<int> decoder_channels, int n_blocks, |
| 66 | bool use_attention, bool use_center) |