Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
trace_model.py 338 Bytes
Newer Older
Edgar Andrés Margffoy Tuay's avatar
PR: Add CMake build and function tracing tests (#2577)  
Edgar Andrés Margffoy Tuay committed
1
2
3
4
5
6
7
8 
import os.path as osp

import torch
import torchvision

HERE = osp.dirname(osp.abspath(__file__))
ASSETS = osp.dirname(osp.dirname(HERE))
Vasilis Vryniotis's avatar
Port Multi-weight support from prototype to main (#5618)  
Vasilis Vryniotis committed
9 
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None, weights_backbone=None)
Edgar Andrés Margffoy Tuay's avatar
PR: Add CMake build and function tracing tests (#2577)  
Edgar Andrés Margffoy Tuay committed
10
11
12
13 
model.eval()

traced_model = torch.jit.script(model)
traced_model.save("fasterrcnn_resnet50_fpn.pt")