Error while trying to save a quantized model using jit

APaul: <p>I followed the steps mentioned in the pytorch <a href="https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html" rel="noopener nofollow ugc">quantization tutorial</a> and tried to save the model using jit but getting an error.</p>
<p>Here is the model structure</p>
<pre><code class="lang-auto">from torch.quantization import QuantStub, DeQuantStub
from loss import TripletLoss

class NN1_BN_FaceNet(nn.Module):

    def __init__(self, classify=False, embedding_size = 128, device=None):

        super(NN1_BN_FaceNet, self).__init__()
        
        self.conv1 = nn.Sequential(nn.BatchNorm2d(3), nn.Conv2d(in_channels=3, out_channels=64, kernel_size=7, stride=2, padding=3), nn.BatchNorm2d(64), nn.ReLU(inplace=True))
        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)
        
        self.conv2a = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, stride=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True))
        self.conv2 = nn.Sequential(nn.Conv2d(64, 192, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(192), nn.ReLU(inplace=True))
        self.pool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=True)

        self.conv3a = nn.Sequential(nn.Conv2d(192, 192, kernel_size=1, stride=1), nn.BatchNorm2d(192), nn.ReLU(inplace=True))
        self.conv3 = nn.Sequential(nn.Conv2d(192, 384, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(384), nn.ReLU(inplace=True))
        self.pool3 = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)

        self.conv4a = nn.Sequential(nn.Conv2d(384, 384, kernel_size=1, stride=1), nn.BatchNorm2d(384), nn.ReLU(inplace=True))
        self.conv4 = nn.Sequential(nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True))
        
        self.conv5a = nn.Sequential(nn.Conv2d(256, 256, kernel_size=1, stride=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True))
        self.conv5 = nn.Sequential(nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True))

        self.conv6a = nn.Sequential(nn.Conv2d(256, 256, kernel_size=1, stride=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True))
        self.conv6 = nn.Sequential(nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True))

        self.pool4 = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)

        self.fc1 = nn.Sequential(nn.Linear(256*7*7, 4*128), nn.ReLU(inplace=True), nn.Dropout())
        self.fc7128 = nn.Sequential(nn.Linear(4*128, embedding_size))
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

        if device is None:
            self.device = torch.device('cpu')
        else:
            self.device = device
            self.to(device)

        def l2_norm(self, input):
            input_size = input.size()
            buffer = torch.pow(input, 2)
            normp = torch.sum(buffer, 1).add_(1e-10)
            norm = torch.sqrt(normp)
            _output = torch.div(input, norm.view(-1, 1).expand_as(input))
            output = _output.view(input_size)
            return output

    def freeze_all(self):
        for param in self.parameters():
            param.requires_grad = False

    def unfreeze_all(self):
        for param in self.parameters():
            param.requires_grad = True

    def freeze_fc(self):
        for param in self.fc.parameters():
            param.requires_grad = False

    def unfreeze_fc(self):
        for param in self.fc.parameters():
            param.requires_grad = True

    def freeze_only(self, freeze):
        for name, child in self.named_children():
            if name in freeze:
                for param in child.parameters():
                    param.requires_grad = False
            else:
                for param in child.parameters():
                    param.requires_grad = True

    def unfreeze_only(self, unfreeze):
        for name, child in self.named_children():
            if name in unfreeze:
                for param in child.parameters():
                    param.requires_grad = True
            else:
                for param in child.parameters():
                    param.requires_grad = False

    def forward(self, x):
        
        x = self.quant(x)
        x = self.conv1(x)
        x = self.pool1(x)
                
        x = self.conv2a(x)
        x = self.conv2(x)
        x = self.pool2(x)

        x = self.conv3a(x)
        x = self.conv3(x)
        x = self.pool3(x)

        x = self.conv4a(x)
        x = self.conv4(x)

        x = self.conv5a(x)
        x = self.conv5(x)

        x = self.conv6a(x)
        x = self.conv6(x)

        x = self.pool4(x)

        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = self.fc7128(x)
        x = nn.functional.normalize(x, p=2, dim=1) 
        x = self.dequant(x)       
        return x
    
    def forward_classifier(self, x):
        features = self.forward(x)
        return features

    def fuse_model(self):
        for m in self.modules():
            if type(m) == nn.Sequential:
                if type(m[0])==nn.BatchNorm2d:
                    self.conv1[0] = nn.Identity()
                    torch.quantization.fuse_modules(self.conv1, ['1', '2', '3'], inplace=True)
                elif type(m[0])==nn.Conv2d and type(m[1])==nn.BatchNorm2d and type(m[2])==nn.ReLU:
                    torch.quantization.fuse_modules(m, ['0', '1', '2'], inplace=True)
                elif (type(m[0])==nn.Linear and len(m)&gt;1):
                    torch.quantization.fuse_modules(m, ['0', '1'], inplace=True)
                else:       
                    print ('No fusion performed on this layer')
                    print(m)
        print('Fusion Complete')
</code></pre>
<p>I am getting the following error when trying to save the model using torch.jit</p>
<pre><code class="lang-auto">torch.jit.save(torch.jit.script(float_model), saved_model_dir + scripted_float_model_file)
</code></pre>
<p>The error message</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-19-d0b6b1651050&gt; in &lt;module&gt;
      6                 columns=[ 'Train_Accuracy', 'Train_Threshold', 'Test_Accuracy', 'Test_Threshold'])
      7 Accuracy_Threshold_Table.to_csv(saved_model_dir+'Acc_Thres_Train_'+str(np.round(np.mean(accuracy_train), 3))+'_Test_'+str(np.round(np.mean(accuracy_test), 3))+'.csv')
----&gt; 8 torch.jit.save(torch.jit.script(float_model), saved_model_dir + scripted_quantized_model_file_S2)
      9 # try:
     10 #     torch.jit.save(torch.jit.script(float_model), saved_model_dir + scripted_quantized_model_file_S2)

~/anaconda3/envs/py39/lib/python3.9/site-packages/torch/jit/_script.py in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1255     if isinstance(obj, torch.nn.Module):
   1256         obj = call_prepare_scriptable_func(obj)
-&gt; 1257         return torch.jit._recursive.create_script_module(
   1258             obj, torch.jit._recursive.infer_methods_to_compile
   1259         )

~/anaconda3/envs/py39/lib/python3.9/site-packages/torch/jit/_recursive.py in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
    449     if not is_tracing:
    450         AttributeTypeIsSupportedChecker().check(nn_module)
--&gt; 451     return create_script_module_impl(nn_module, concrete_type, stubs_fn)
    452 
    453 def create_script_module_impl(nn_module, concrete_type, stubs_fn):

~/anaconda3/envs/py39/lib/python3.9/site-packages/torch/jit/_recursive.py in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    515     # Compile methods if necessary
    516     if concrete_type not in concrete_type_store.methods_compiled:
--&gt; 517         create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    518         # Create hooks after methods to ensure no name collisions between hooks and methods.
    519         # If done before, hooks can overshadow methods that aren't exported.
~/anaconda3/envs/py39/lib/python3.9/site-packages/torch/jit/_recursive.py in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    366     property_rcbs = [p.resolution_callback for p in property_stubs]
    367 
--&gt; 368     concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
    369 
    370 def create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs):

RuntimeError: 

normalize(Tensor input, float p=2., int dim=1, float eps=9.9999999999999998e-13, Tensor? out=None) -&gt; (Tensor):
Expected a value of type 'float' for argument 'p' but instead found type 'int'.
:
  File "&lt;ipython-input-4-129ce50ca4a2&gt;", line 145
        
        x = self.dequant(x)
        x = nn.functional.normalize(x, p=2, dim=1)                
            ~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
        return x
</code></pre>
<p>Any suggestions would be appreciated</p>