Quantized Squeeze block MobilenetV3

Hi, I’m trying to quantize mobilenet V3 but get stuck in quantizing Squeeze Block. When I train it in subset dataset and use convert to quantized model, it works. but when I evaluate quantized model in eval set, It throws an error of Mul operation. Can you tell me how can I implement Squeeze block (SElayer). thank you so much.

class SqueezeBlock(nn.Module):
    def __init__(self, exp_size, divide=4):
        super(SqueezeBlock, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.dense = nn.Sequential(
            nn.Linear(exp_size, exp_size // divide),
            nn.ReLU(inplace=False),
            nn.Linear(exp_size // divide, exp_size),
            h_sigmoid()
        )
        self.mul = torch.nn.quantized.FloatFunctional()

    def forward(self, x):
        batch, channels, height, width = x.size()
        out = self.avg_pool(x).view(batch, channels)
        out = self.dense(out)
        out = out.view(batch, channels, 1, 1)
        return self.mul.mul(x,out)
File "/home/X/Documents/thancuong/mb3_quantized/quantized_mb3.py", line 121, in forward
    return self.mul.mul(x,out)
  File "/home/X/.virtualenvs/torch_0.4/lib/python3.5/site-packages/torch/nn/quantized/modules/functional_modules.py", line 146, in mul
    zero_point=self.zero_point)
RuntimeError: Mul operands must be the same size! (check_inputs at /pytorch/aten/src/ATen/native/quantized/cpu/qmul.cpp:20)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x33 (0x7f077b622813 in /home/X/.virtualenvs/torch_0.4/lib/python3.5/site-packages/torch/lib/libc10.so)

You can actually try to comment out the two lines as https://github.com/pytorch/pytorch/pull/30442, since the tensor iterator supports broadcast.

thank you, I will try it