Question about quantized model save & load

I have a model like:

model = MyModel()

and I quantized part of it:

model.head = quantize(model.head)

and it runs fine:

output = model(input)

but when I save and load it like:

torch.save(model, 'model.pt')
new_m = torch.load('model.pt')

an error throwed:
AttributeError: ‘ConvReLU2d’ object has no attribute ‘_modules’

How can I fix this?

this might mean ConvReLU2d module does not call super().__init__() in its constructor

can you show the complete quantization code?

Thanks for reply.

Well I cant’ show you all the quantization code since we built a toolkit that based torch fx graph, and it’s kind complicated.

The key point is that we use symbolic_trace to retrieve the model graph from eager model, then use torch.ao.quantization.fx.fuse to fuse batch norm and other operators, then lower it to GPU.

The weird thing is that the ConvReLU2d we are using is imported from here , and it does called super().__init__(), so I don’t think that’s the root cause.

I see, I did see this error before while I was working on fx/eager, but can’t recall how these are resolved anymore..

If it’s possible to update your tools, I’d recommend to migrate to the pt2e quant toolchain: Quantization — PyTorch main documentation it also uses fx graph and the graph capture is better than symbolic_trace, more details can be found in (prototype) PyTorch 2 Export Post Training Quantization — PyTorch Tutorials 2.6.0+cu124 documentation