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)