romain-keramitas-prl

import numpy as np
from onnxruntime import InferenceSession
import torch
import torch.nn as nn

# Create simple model with one Gather node
model = nn.Embedding(num_embeddings=10, embedding_dim=5)
x = torch.randint(0, 10, (3,))
torch.onnx.export(
    model,

romain-keramitas-prl

import os

import numpy as np
from onnx import TensorProto, helper, save
from onnxruntime import InferenceSession

hidden_size = 10
num_heads = 2
head_size = 5