From cv2 to TORCH:
 cv::cvtColor(frame, frame, CV_BGR2RGB);
 frame.convertTo(frame, CV_32FC3, 1.0f / 255.0f);
 auto input_tensor = torch::from_blob(frame.data, {1, frame_h, frame_w, kCHANNELS});
 input_tensor = input_tensor.permute({0, 3, 1, 2});

From TORCh to cv2:
input_tensor = input_tensor.to(at::kCUDA);
                torch::Tensor out_tensor = module->forward({input_tensor}).toTensor();

                out_tensor = out_tensor.squeeze().detach().permute({1, 2, 0});
                out_tensor = out_tensor.mul(255).clamp(0, 255).to(torch::kU8);
                out_tensor = out_tensor.to(torch::kCPU);
                cv::Mat resultImg(frame_h, frame_w, CV_8UC3);
                std::memcpy((void *) resultImg.data, out_tensor.data_ptr(), sizeof(torch::kU8) * out_tensor.numel());