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.