Quantizing MaxPool3D

Hello. I am trying to employ post-training quantization on a SqueezeNet 3D and am getting the following error:

RuntimeError: Could not run ‘aten::max_pool3d_with_indices’ with arguments from the ‘QuantizedCPU’ backend. ‘aten::max_pool3d_with_indices’ is only available for these backends: [CPU, CUDA, Named, Autograd, Profiler, Tracer].

I checked the list of supported operators (Quantization API Reference — PyTorch 2.1 documentation) and it seems that MaxPool3D is not supported. How can I circumvent this issue?

if an operator does not have a quantized kernel, you can run it in fp32. For example,

class M(torch.nn.Module):
  def __init__(...):
    ...
    self.dequant_1 = torch.quantization.DeQuantStub()
    self.quant_1 = torch.quantization.QuantStub()
    ...

  def forward(self, ...):
    ...
    # (computation in int8)
    ...
    # convert from int8 to fp32
    x_1 = self.dequant_1(x_0)
    # run in fp32
    x_2 = self.maxpool(x_1)
    # convert back to int8, if needed
    x_3 = self.quant_q(x_2)
    ...
    
  

Thank you very much for your reply. That actually answers my question, but hopefully, I can take the opportunity to ask you another one. I am using a pretrained SqueezeNet 3D and in order to be able to load saved models I need to keep the nn.Sequential modules (namely self.features and self.classifier). I tried doing something like this:

        if self.quantize:
            # run Pooling operations and dropout in FP32
            x = self.quant(x)
            # x = self.features(x)
            for m in self.features:
                if isinstance(m, nn.MaxPool3d):
                    x = self.dequant(x)
                    x = m(x)
                    x = self.quant(x)
                else:
                    x = m(x)
            
            for c in self.classifier:
                if isinstance(c, nn.AvgPool3d) or isinstance(c, nn.Dropout):
                    x = self.dequant(x)
                    x = c(x)
                    x = self.quant(x)
                else:
                    x = c(x)
            x = x.view(x.size(0), -1)
            x = self.dequant(x)
            return x
        else:
            x = self.features(x)
            x = self.classifier(x)
            return x.view(x.size(0), -1)

This gives the following error:

RuntimeError:
Cannot re-assign ‘c’ because it has type value of type ‘torch.torch.nn.modules.pooling.AvgPool3d’ and c is not a first-class value. Only reassignments to first-class values are allowed:
File “/home/ctm/afonso/easyride/acceleration/src/models/squeezenet.py”, line 152

        for c in self.classifier:
            if isinstance(c, nn.AvgPool3d) or isinstance(c, nn.Dropout):
               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                x = self.dequant(x)
                x = c(x)

Although not understanding why, I actually managed to fix the problem using the following code extract instead of the classifier loop:

            x = self.dequant(x)
            x = self.classifier[0](x)
            x = self.quant(x)
            x = self.classifier[1](x)
            x = self.classifier[2](x)
            x = self.dequant(x)
            x = self.classifier[-1](x)

I still do not fully understand how the TorchScript inspector works. I would appreciate if you could give me some insights on this behaviour.