How to quantify BatchNorm2d layer？ about concrete-ml HOT 9 CLOSED

The following tensors were expected to be quantized, but the values found during calibration do not appear to be quantized.

result.add_module('conv',qnn.QuantConv2d(in_channels=in_channels,out_channels=out_channels, kernel_size=kernel_size, stride=stride,padding=padding, weight_bit_width=bit, weight_quant=weight_quant, return_quant_tensor=True, bias=False))
result.add_module('bn', nn.BatchNorm2d(num_features=out_channels, eps=1e-4))

import brevitas
import brevitas.nn as qnn
import torch
import torch.nn as nn
from brevitas.quant import Int8ActPerTensorFloat, Int8WeightPerTensorFloat

def conv_bn(in_channels, out_channels, kernel_size, stride, padding,bit,
                 act_quant: brevitas.quant = Int8ActPerTensorFloat,
                 weight_quant: brevitas.quant = Int8WeightPerTensorFloat,):
    result = nn.Sequential()
    result.add_module('conv', qnn.QuantConv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
                                padding=padding, weight_bit_width=bit, weight_quant=weight_quant, return_quant_tensor=True, bias=False))
    result.add_module('bn', nn.BatchNorm2d(num_features=out_channels, eps=1e-4))
    return result

class QuantmodelBlock(nn.Module):
    def __init__(self, in_channels, out_channels ,stride: int,kernel_size: int, padding: int, bit: int,
                 act_quant: brevitas.quant = Int8ActPerTensorFloat,
                 weight_quant: brevitas.quant = Int8WeightPerTensorFloat,):
        super(QuantmodelBlock, self).__init__()
        self.in_channels = in_channels
        self.bit = bit
        padding_11 = padding - kernel_size // 2
        self.nonlinearity = qnn.QuantReLU(return_quant_tensor=True, bit_width=bit, act_quant=act_quant)
        self.rbr_identity = nn.BatchNorm2d(num_features=in_channels,eps=1e-4) if out_channels == in_channels and stride == 1 else None
        self.rbr_dense = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,padding=padding, bit=bit)
        print(in_channels, out_channels, kernel_size, stride, padding)
        self.rbr_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride,padding=padding_11, bit=bit)

    def forward(self, inputs):
        if self.rbr_identity is None:
            id_out = 0
        else:
            id_out = self.rbr_identity(inputs)
        return self.nonlinearity(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out).value

class Quantmodel(nn.Module):

    def __init__(self, num_blocks, width_multiplier,
            bit: int,
            output_size: int,
            act_quant: brevitas.quant = Int8ActPerTensorFloat,
            weight_quant: brevitas.quant = Int8WeightPerTensorFloat,):
        super(Quantmodel, self).__init__()
        assert len(width_multiplier) == 4
        self.in_planes = min(16, int(16 * width_multiplier[0]))
        self.stage0 = QuantmodelBlock(in_channels=3, out_channels=self.in_planes, kernel_size=3, padding=1, bit=bit, stride=1)
        self.stage1 = self._make_stage(int(16 * width_multiplier[0]), num_blocks[0], stride=1, bit = bit)
        self.stage2 = self._make_stage(int(32 * width_multiplier[1]), num_blocks[1], stride=1 ,bit = bit)
        self.stage3 = self._make_stage(int(64 * width_multiplier[2]), num_blocks[2], stride=2, bit = bit)
        self.stage4 = self._make_stage(int(64 * width_multiplier[3]), num_blocks[3], stride=2, bit = bit)
        self.linear1 = qnn.QuantLinear(
            in_features=int(64 * width_multiplier[3]),
            out_features=output_size,
            weight_quant=weight_quant,
            weight_bit_width=bit,
            bias=True,
            return_quant_tensor=True,
        )
        self.identity = qnn.QuantIdentity(
             bit_width=bit, act_quant=act_quant, return_quant_tensor=True
    )
        self.identity2 = qnn.QuantIdentity(
            bit_width=bit, act_quant=act_quant, return_quant_tensor=True
        )
        self.identity3 = qnn.QuantIdentity(
            bit_width=bit, act_quant=act_quant, return_quant_tensor=True
        )
        self.gap = nn.AvgPool2d(kernel_size=8, stride=1, padding=0, ceil_mode=False)

    def _make_stage(self, planes, num_blocks, stride,bit):
        strides = [stride] + [1]*(num_blocks-1)
        blocks = []
        for stride in strides:
            blocks.append(QuantmodelBlock(in_channels=self.in_planes, out_channels=planes, stride=stride,bit=bit,kernel_size=3,padding=1))
            self.in_planes = planes
        return nn.ModuleList(blocks)

    def forward(self, x):
        out = self.identity(x)
        out = self.stage0(out)
        for stage in (self.stage1, self.stage2, self.stage3, self.stage4):
            for block in stage:
                out = block(out)
        out = self.gap(out)
        out = self.identity2(out)
        out = torch.flatten(out, 1)
        out = self.identity3(out)
        out = self.linear1(out)
        return out.value

def create_Quantmodel(bit: int, output_size: int):
    return Quantmodel(num_blocks=[4, 4, 4, 1],width_multiplier=[1, 1, 1, 1], bit=bit, output_size=output_size)

        return self.nonlinearity(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out).value

def __init__():
  ....
  self.out_quant = qnn.QuantIdentity()

def forward(..):
  ....
        return self.nonlinearity(self.out_quant(self.rbr_dense(inputs)) + self.out_quant(self.rbr_1x1(inputs)) + self.out_quant(id_out)).value