Skip to content

Instantly share code, notes, and snippets.

@steerapi

steerapi/torchjs_export.py

Created November 6, 2018 22:28
Show Gist options
  • Save steerapi/114a33820146de68848e43b18016d8af to your computer and use it in GitHub Desktop.
Save steerapi/114a33820146de68848e43b18016d8af to your computer and use it in GitHub Desktop.
Download ZIP
Raw
torchjs_export.py
import torch
import torchvision
# An instance of your model.
model = torchvision.models.resnet18(pretrained=True)
# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)
# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("resnet18.pt")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment