Any way to export Densenet with memoryefficient=true to onnx?

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)

The memory_efficient=True option uses internally torch.utils.checkpoint to trade compute for memory.
This operator is not defined in ONNX, which raises this error.
However, if you don’t necessarily need checkpointing, you could just set this argument to False and the export should work.