Quantization of full models

[Wheest]

I am trying to quantize some simple CIFAR10 models. I am following the documentation on quantization.
You can see this gist which I am using, which I have removed any pretraining from to keep simple, which reproduces my issue.
Note that there is a dependency of the pytorch-lighting-cifar pacakge (pip3 install pytorch-lighting-cifar) which has the model architectures.

Trying an example with python quantize_standalone.py --model_arch mobilenetv2 --epochs 2 fails with:

Traceback (most recent call last):
  File "quantize_standalone.py", line 150, in <module>
    main(args)
  File "quantize_standalone.py", line 117, in main
    tsmodel = model.to_torchscript()
  File "/home/user/.local/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/user/.local/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1926, in to_torchscript
    torchscript_module = torch.jit.script(self.eval(), **kwargs)
  File "/home/user/.local/lib/python3.8/site-packages/torch/jit/_script.py", line 1265, in script
    return torch.jit._recursive.create_script_module(
  File "/home/user/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 454, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/user/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 520, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/home/user/.local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 371, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError:
undefined value F:
  File "/home/user/.local/lib/python3.8/site-packages/pytorch_lightning/callbacks/quantization.py", line 99
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
              ~ <--- HERE
        out = self.layers(out)
        out = F.relu(self.bn2(self.conv2(out)))

Since it is a callback, out = F.relu(self.bn1(self.conv1(x))) is not actually in this quantization.py, so I am not sure why it fails.
However, if I add the line to import torch.nn.functional as F at the start of the file quantization.py, I get further, however then I fail with:

Traceback (most recent call last):
  File "quantize_standalone.py", line 150, in <module>
    main(args)
  File "quantize_standalone.py", line 118, in main
    y2 = tsmodel(tsmodel.quant(batch))
  File "/home/user/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "/home/user/.local/lib/python3.8/site-packages/pytorch_lightning/callbacks/quantization.py", line 99, in forward
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
                     ~~~~~~~~ <--- HERE
        out = self.layers(out)
        out = F.relu(self.bn2(self.conv2(out)))
  File "/home/user/.local/lib/python3.8/site-packages/torch/nn/quantized/modules/batchnorm.py", line 17, in forward
    def forward(self, input):
        return torch.ops.quantized.batch_norm2d(input, self.weight, self.bias, self.running_mean,
               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                                                self.running_var, self.eps, self.scale, self.zero_point)
RuntimeError: Cannot input a tensor of dimension other than 0 as a scalar argument

I get the same issue running other models such as resnet18 and mobilenetv1.
Is there something I'm missing? I believe I'm following the docs correctly.
@Borda @awaelchli git blame says you have both committed heavily to the quantization code, is this something you have seen before, or can give any tips to tackle?