用torch写一个topk+add的小模型
import torch
print(torch.__version__)
class Net1(torch.nn.Module):
def __init__(self):
super(Net1, self).__init__()
def forward(self, x):
output = torch.topk(x,1000)
print("**************",type(output),"********")
input_add = torch.randint(1,10000,(1,1000)).to(torch.int64)
output_1 = torch.add(input=output[0],other=input_add)
return output_1
model = Net1()
dummy_input = torch.randint(1, 10000,(1,10336))
dummy_input = dummy_input.to(torch.int64)
input_names = ["input_1"]
output_names = ["output_1"]
torch.onnx.export(model, dummy_input, "topk.onnx",input_names=input_names,output_names=output_names,opset_version=12)
![在这里插入图片描述](https://i-blog.csdnimg.cn/blog_migrate/d7857547d95a47db4f65065d670bdd41.png)