Hello
I’m trying to do QAT -> Torchscript but am getting an error.
My model is

 Click here 
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.quantization import QuantStub, DeQuantStub

def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, activation_fn=nn.ReLU):
    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
    """
    return nn.Sequential(
        nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
                  groups=in_channels, stride=stride, padding=padding),
        activation_fn(),
        nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)
    )


class SSD(nn.Module):
    def __init__(self, num_classes: int, is_test=False, config=None, device=None, activation_fn=nn.ReLU):
        """Compose a SSD model using the given components.
        """
        super(SSD, self).__init__()
        self.base_channel = 16
        self.num_classes = num_classes
        self.is_test = is_test

        if device:
            self.device = device
        else:
            self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

        if config:
            self.center_variance = torch.tensor([config['center_variance']], device=device)
            self.size_variance = torch.tensor([config['size_variance']], device=device)
            self.priors = config['priors'].to(self.device)
        else:
            self.center_variance = torch.tensor([0.1], device=device)
            self.size_variance = torch.tensor([0.2], device=device)
        
        self.extras = nn.Sequential(
            nn.Conv2d(in_channels=self.base_channel * 16, out_channels=self.base_channel * 4, kernel_size=1),
            activation_fn(),
            SeperableConv2d(in_channels=self.base_channel * 4, out_channels=self.base_channel *
                               16, kernel_size=3, stride=2, padding=1, activation_fn=activation_fn),
            activation_fn()
        )

        self.regression_headers0 = SeperableConv2d(in_channels=self.base_channel * 4, out_channels=3 *
                                                   4, kernel_size=3, padding=1, activation_fn=activation_fn)
        self.regression_headers1 = SeperableConv2d(in_channels=self.base_channel * 8, out_channels=2 *
                                                   4, kernel_size=3, padding=1, activation_fn=activation_fn)
        self.regression_headers2 = SeperableConv2d(in_channels=self.base_channel * 16, out_channels=2 *
                                                   4, kernel_size=3, padding=1, activation_fn=activation_fn)
        self.regression_headers3 = nn.Conv2d(in_channels=self.base_channel * 16,
                                             out_channels=3 * 4, kernel_size=3, padding=1)

        self.classification_headers0 = SeperableConv2d(in_channels=self.base_channel * 4, out_channels=3 *
                                                       num_classes, kernel_size=3, padding=1, activation_fn=activation_fn)
        self.classification_headers1 = SeperableConv2d(in_channels=self.base_channel * 8, out_channels=2 *
                                                       num_classes, kernel_size=3, padding=1, activation_fn=activation_fn)
        self.classification_headers2 = SeperableConv2d(in_channels=self.base_channel * 16, out_channels=2 *
                                                       num_classes, kernel_size=3, padding=1, activation_fn=activation_fn)
        self.classification_headers3 = nn.Conv2d(in_channels=self.base_channel *
                                              16, out_channels=3 * num_classes, kernel_size=3, padding=1)

        def conv_bn(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                nn.BatchNorm2d(oup),
                activation_fn()
            )

        def conv_dw(inp, oup, stride):
            return nn.Sequential(
                nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                nn.BatchNorm2d(inp),
                activation_fn(),

                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                nn.BatchNorm2d(oup),
                activation_fn(),
            )
        self.backbone_chunk1 = nn.Sequential(
            conv_bn(3, self.base_channel, 2),  # 160*120
            conv_dw(self.base_channel, self.base_channel * 2, 1),
            conv_dw(self.base_channel * 2, self.base_channel * 2, 2),  # 80*60
            conv_dw(self.base_channel * 2, self.base_channel * 2, 1),
            conv_dw(self.base_channel * 2, self.base_channel * 4, 2),  # 40*30
            conv_dw(self.base_channel * 4, self.base_channel * 4, 1),
            conv_dw(self.base_channel * 4, self.base_channel * 4, 1),
            # BasicRFB(self.base_channel * 4, self.base_channel * 4, stride=1, scale=1.0, activation_fn=activation_fn)
        )
        self.backbone_chunk2 = nn.Sequential(
            conv_dw(self.base_channel * 4, self.base_channel * 8, 2),  # 20*15
            conv_dw(self.base_channel * 8, self.base_channel * 8, 1),
            conv_dw(self.base_channel * 8, self.base_channel * 8, 1),
        )
        self.backbone_chunk3 = nn.Sequential(
            conv_dw(self.base_channel * 8, self.base_channel * 16, 2),  # 10*8
            conv_dw(self.base_channel * 16, self.base_channel * 16, 1)
        )

        self.quant0 = QuantStub()
        self.quant1 = QuantStub()
        self.quant2 = QuantStub()
        self.quant3 = QuantStub()
        self.dequant = DeQuantStub()


    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        confidences = []
        locations = []

        x = self.quant0(x)
        for layer in self.backbone_chunk1:
            x = layer(x)
        x = self.dequant(x)
        confidence = self.classification_headers0(x)
        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
        location = self.regression_headers0(x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)
        confidences.append(confidence)
        locations.append(location)

        # x = self.quant1(x)
        for layer in self.backbone_chunk2:
            x = layer(x)
        # x = self.dequant(x)
        confidence = self.classification_headers1(x)
        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
        location = self.regression_headers1(x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)
        confidences.append(confidence)
        locations.append(location)

        # x = self.quant2(x)
        for layer in self.backbone_chunk3:
            x = layer(x)
        # x = self.dequant(x)
        confidence = self.classification_headers2(x)
        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
        location = self.regression_headers2(x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)
        confidences.append(confidence)
        locations.append(location)

        # x = self.quant3(x)
        x = self.extras(x)
        # x = self.dequant(x)
        confidence = self.classification_headers3(x)
        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
        location = self.regression_headers3(x)
        location = location.permute(0, 2, 3, 1).contiguous()
        location = location.view(location.size(0), -1, 4)
        confidences.append(confidence)
        locations.append(location)

        confidences = torch.cat(confidences, 1)
        locations = torch.cat(locations, 1)

        return confidences, locations

    def load(self, model):
        self.load_state_dict(torch.load(model, map_location=lambda storage, loc: storage))

    def save(self, model_path):
        torch.save(self.state_dict(), model_path)



I do QAT -> torchscript and test it by running the code:
from ssd import SSD
...
net = SSD(num_classes=2, device=device, config=config)
...
net.load(trained_model_path)
net.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(net, inplace=True)

for epoch in range(num_epochs):
    train() etc...

    net.eval()
    net.cpu()
    # convert to quantised
    quant_net = deepcopy(net)
    quant_net = torch.quantization.convert(quant_net, inplace=False)
    quant_net.save(os.path.join(args.checkpoint_folder,
                                          f"quantised-net.pth"))
            
    m = torch.jit.script(quant_net)
    m.cpu()
    dummy = torch.randn(1, 3, 480, 640).cpu().float()
    a = m.forward(dummy) # test to see if scripted module works

    torch.jit.save(m, os.path.join(args.checkpoint_folder, f"jit-net.pt"))

    net.to(DEVICE)

and I get the error on the line  a = m.forward(dummy) :
Traceback (most recent call last):
  File "train_testt.py", line 360, in <module>
    a = m.forward(dummy)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "/home/joel/Desktop/Ultra-Light-Fast-Generic-Face-Detector-1MB/minimod.py", line 124, in forward
            x = layer(x)
        x = self.dequant(x)
        confidence = self.classification_headers0(x)
                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        confidence = confidence.permute(0, 2, 3, 1).contiguous()
        confidence = confidence.view(confidence.size(0), -1, self.num_classes)
  File "/home/joel/anaconda3/envs/nightlytorch/lib/python3.8/site-packages/torch/nn/modules/container.py", line 117, in forward
    def forward(self, input):
        for module in self:
            input = module(input)
                    ~~~~~~ <--- HERE
        return input
  File "/home/joel/anaconda3/envs/nightlytorch/lib/python3.8/site-packages/torch/nn/modules/container.py", line 117, in forward
    def forward(self, input):
        for module in self:
            input = module(input)
                    ~~~~~~ <--- HERE
        return input
  File "/home/joel/anaconda3/envs/nightlytorch/lib/python3.8/site-packages/torch/nn/quantized/modules/conv.py", line 326, in forward
        if len(input.shape) != 4:
            raise ValueError("Input shape must be `(N, C, H, W)`!")
        return ops.quantized.conv2d(
               ~~~~~~~~~~~~~~~~~~~~ <--- HERE
            input, self._packed_params, self.scale, self.zero_point)
RuntimeError: Could not run 'quantized::conv2d.new' with arguments from the 'CPU' backend. 'quantized::conv2d.new' is only available for these backends: [QuantizedCPU, BackendSelect, Named, Autograd, Profiler, Tracer, Autocast, Batched].

QuantizedCPU: registered at /opt/conda/conda-bld/pytorch_1594145889316/work/aten/src/ATen/native/quantized/cpu/qconv.cpp:736 [kernel]
BackendSelect: fallthrough registered at /opt/conda/conda-bld/pytorch_1594145889316/work/aten/src/ATen/core/BackendSelectFallbackKernel.cpp:3 [backend fallback]
Named: registered at /opt/conda/conda-bld/pytorch_1594145889316/work/aten/src/ATen/core/NamedRegistrations.cpp:7 [backend fallback]
Autograd: fallthrough registered at /opt/conda/conda-bld/pytorch_1594145889316/work/aten/src/ATen/core/VariableFallbackKernel.cpp:31 [backend fallback]
Profiler: registered at /opt/conda/conda-bld/pytorch_1594145889316/work/torch/csrc/autograd/profiler.cpp:677 [backend fallback]
Tracer: fallthrough registered at /opt/conda/conda-bld/pytorch_1594145889316/work/torch/csrc/jit/frontend/tracer.cpp:960 [backend fallback]
Autocast: fallthrough registered at /opt/conda/conda-bld/pytorch_1594145889316/work/aten/src/ATen/autocast_mode.cpp:375 [backend fallback]
Batched: registered at /opt/conda/conda-bld/pytorch_1594145889316/work/aten/src/ATen/BatchingRegistrations.cpp:149 [backend fallback]

This error does not occur if I remove all QAT/quantization lines and just jit.script the original model. The same error still occurs if I remove all Quant/DeQuantStubs in the model.
Does anyone know why this error occurs?
Could I also ask whether the commented out QuantStubs/DeStubs in the model are correctly placed?
Thank you!