Hi,
I am trying to export Densenet121 to ONNX format, but I am stuck with
the process, getting exception :
ONNX export failed: Couldn’t export Python operator CheckpointFunction
Is there any way to go around the exception?
My code for export is:
def export_model(model):
sample_batch_size = 1
channel = 3
height = 224
width = 224
dummy_input = torch.randn(sample_batch_size, channel, height, width)
torch.onnx.export(model, dummy_input, "onnx_model_name.onnx", input_names=['input'], output_names=['output'])
model = torch.hub.load('pytorch/vision', 'densenet121', pretrained=True, memory_efficient=True)
export_model(model)