Skip to content

Instantly share code, notes, and snippets.

@jaemin93

jaemin93/exporting_pytorch_to_onnx.py

Created June 30, 2020 08:17
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save jaemin93/c4ff91eb9275bce2aef2da6a6a519d3e to your computer and use it in GitHub Desktop.
Save jaemin93/c4ff91eb9275bce2aef2da6a6a519d3e to your computer and use it in GitHub Desktop.
Download ZIP
Raw
exporting_pytorch_to_onnx.py
import torch
from classification_model import Net
from torch.onnx import OperatorExportTypes
model = Net()
model.eval()
x = torch.zeros([1, 1, 28, 28])
print(x.shape)
torch.onnx.export(model, x, "test.onnx", verbose=True, operator_export_type=OperatorExportTypes.ONNX)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment