This simple script will illustrate what I’m trying to do
import torch
from torch import nn
from torch.quantization.quantize_fx import prepare_fx

class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        B = x.shape[0]
        B = B * 1
        return x

model = Model()
model.eval()

qconfig = torch.quantization.get_default_qconfig("fbgemm")
model = prepare_fx(model, {"": qconfig})

x = torch.arange(4)

print(model(x))

Right now it raises AttributeError: 'int' object has no attribute 'numel' because of the B * 1 line.
By the way, the real example for what I’m trying to quantize is here 1.