import torch torch.onnx.export( final_model, torch.randn(1, 3, 256, 256), 'models/hot_dog_model_resnet18_256_256.onnx', do_constant_folding=True, export_params=True, input_names=['image_1_3_256_256'], output_names=['hot_dog'], opset_version=11 )