How to convert a model to ONNX before conversion?

quantization
1

For some reason, I need to store a intermediate state model into ONNX，The state of this model is between ‘prepare’ and ‘convert’，My current approach is as follows

import torch
import torch.nn as nn
#from torch.ao.quantization import QConfigMapping
from torch.quantization.quantize_fx import prepare_fx, convert_fx, prepare_qat_fx

class Net(torch.nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.bn1 = nn.BatchNorm2d(6)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(6, 1, 1)
        self.bn2 = nn.BatchNorm2d(1)
        self.relu2 = nn.ReLU()
    def forward(self, x):
        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.conv2(x)))
        return x

example_inputs = torch.randn((1, 3, 48, 48))
model = Net()
model.train()
qconfig_dict = {"":torch.quantization.get_default_qat_qconfig('qnnpack'),}
prepared = prepare_qat_fx(model,qconfig_dict)
quantized_model = convert_fx(prepared)
torch.onnx.export(prepared, example_inputs, "quantized_model.onnx", input_names=['input1'], output_names=['result'],opset_version=14, export_params=True,  operator_export_type=torch.onnx.OperatorExportTypes.ONNX_FALLTHROUGH, keep_initializers_as_inputs=True )

but i got

File “/usr/local/anaconda3/envs/yolo/lib/python3.8/site-packages/torch/onnx/utils.py”, line 1111, in _model_to_graph
graph, params, torch_out, module = _create_jit_graph(model, args)
File “/usr/local/anaconda3/envs/yolo/lib/python3.8/site-packages/torch/onnx/utils.py”, line 987, in _create_jit_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
File “/usr/local/anaconda3/envs/yolo/lib/python3.8/site-packages/torch/onnx/utils.py”, line 883, in _trace_and_get_graph_from_model
orig_state_dict_keys = torch.jit._unique_state_dict(model).keys()
File “/usr/local/anaconda3/envs/yolo/lib/python3.8/site-packages/torch/jit/_trace.py”, line 71, in _unique_state_dict
filtered_dict[k] = v.detach()
AttributeError: ‘torch.qscheme’ object has no attribute ‘detach’

My PyTorch version is 1.13.1+cu117

Could you please advise on how to resolve this issue? :melting_face: