Method ConvReluBnImpl

lesson2-TensorOperations/main.cpp:22–26

TORCH_MODULE(ConvReluBn);

ConvReluBnImpl::ConvReluBnImpl(int input_channel, int output_channel, int kernel_size) {
    conv = register_module("conv", torch::nn::Conv2d(torch::nn::Conv2dOptions(input_channel, output_channel, kernel_size).padding(1)));
    bn = register_module("bn", torch::nn::BatchNorm2d(output_channel));

}

torch::Tensor ConvReluBnImpl::forward(torch::Tensor x) {
    x = torch::relu(conv->forward(x));