File size: 315 Bytes
6f23931 |
1 2 3 4 5 6 7 8 9 10 11 12 |
from transformers import AutoTokenizer, AutoModel
import torch
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model.eval()
sample = torch.randint(low=0, high=1, size=(1,128))
input = (sample, sample, sample)
torch.onnx.export(model, input, 'pytorch_model.onnx', export_params=True)
|