I am not sure how can this module help in setting custom number of bits, is it possible using quantization module ?.
Also can you please help me with a code snippet to change number of bits for the activation ?  Currently i am using this one but it doesn’t seem to work.
arr1 = [9,8,5,5,7]     # activations quantized 
#arr1 = [2,2,2,2,2]

def quantizee(self,input,output):
    #print("Hello")
    #print('Inside ' + self.__class__.__name__ + ' forward')
    #print('input: ', type(input))
    #print('output: ', type(output))
    
    print('output size:', output.data.size())
    print(arr1[quantizee.counter])
    #out = output.detach()
    #print('out: ', type(out))
    output = quant.min_max_quantize(output.data,arr1[quantizee.counter])
    quantizee.counter += 1
    if (quantizee.counter == 5):
        quantizee.counter = 0



quantizee.counter = 0

count1 = 0          # number of relus you want to quantize count 1 is that number
for name, module in model.named_modules():
    if (isinstance (module, torch.nn.modules.Conv2d) & (count1 < 5)):
        module.register_forward_hook(quantizee)
        count1 += 1'''

i have tried this one too but no success.
#arr1 = [9,8,5,5,7]
arr1 = [2,2,2,2,2]
count1 = 0
l = OrderedDict()
for k, v in state_dict_quant.items():
    if isinstance(v, (nn.ReLU)):
        quant_layer = quant.NormalQuant('{}_quant'.format(k), bits=arr1[count1], quant_func=log_minmax_quantize)
        l['{}_{}_quant'.format(k, type)] = quant_layer
        count1 += 1
        if (count1 == 5):
            count1= 0
    
    else:
        l[k]= v