mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
8c8bd79f32
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/32316 ### Summary Since the Custom Build has been released in 1.4.0, it's time setup CI. To do that, we need 1. Add a python script to generate the yaml file 2. Add new build scripts to circle CI (arm64 only). ### Test Plan - Don't break the current iOS CIs Test Plan: Imported from OSS Differential Revision: D19437362 Pulled By: xta0 fbshipit-source-id: 395e27a582c43663af88d11b1ef974a4687e672c
12 lines
333 B
Python
12 lines
333 B
Python
import torch
|
|
import torchvision
|
|
import yaml
|
|
|
|
model = torchvision.models.mobilenet_v2(pretrained=True)
|
|
model.eval()
|
|
example = torch.rand(1, 3, 224, 224)
|
|
traced_script_module = torch.jit.trace(model, example)
|
|
ops = torch.jit.export_opnames(traced_script_module)
|
|
with open('mobilenetv2.yaml', 'w') as output:
|
|
yaml.dump(ops, output)
|