Export onnx and channel_last

i am using torch.onnx.export to export an onnx model, and it works. but my depolyment enviroment need a channel last model. i can only export channel first model now. Is there some ways to get a channel last model ? thanks for any replies !

Does this work

model = model.to(memory_format=torch.channels_last) and then run the export?