Sign in
android
/
platform
/
external
/
pytorch
/
d8c3b2b10c
/
.
/
ios
/
TestApp
/
benchmark
/
trace_model.py
blob: 2b96b162ef7e8bb79605e6759c049b6f08b2fc4b [
file
] [
log
] [
blame
]
import
torch
import
torchvision
model
=
torchvision
.
models
.
mobilenet_v2
(
pretrained
=
True
)
model
.
eval
()
example
=
torch
.
rand
(
1
,
3
,
224
,
224
)
traced_script_module
=
torch
.
jit
.
trace
(
model
,
example
)
traced_script_module
.
save
(
"model.pt"
)