Quantization multiple branches

Hi, I want to quantize a CNN with two fully connected branches at the end of the network. I don’t know if I should use one or two instances of DequantStub? Both V1 and V2 seem to work but what is the best practice?

class QuantizedModelV1(Module):

    def __init__(self, fp32_model):
        super(QuantizedModel, self).__init__()
        self.model_fp32 = fp32_model
        self.quant = QuantStub()
        self.dequant_b1 = DeQuantStub()
        self.dequant_b2 = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        b1, b2 = self.model_fp32(x)
        b1 = self.dequant_b1(b1)
        b2 = self.dequant_b2(b2)
        return b1, b2
class QuantizedModelV2(Module):

    def __init__(self, fp32_model):
        super(QuantizedModel, self).__init__()
        self.model_fp32 = fp32_model
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        b1, b2 = self.model_fp32(x)
        b1 = self.dequant(b1)
        b2 = self.dequant(b2)
        return b1, b2

DeQuantStub doesn’t carry any state so both are fine. I’d recommend the first one though, just to be more robust.