From 88086ae2a89b470700509e7e2b26af6cf374ae91 Mon Sep 17 00:00:00 2001
From: Jael Gu <mengjia.gu@zilliz.com>
Date: Thu, 15 Dec 2022 14:14:55 +0800
Subject: [PATCH] Fix run.py for new op.model

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
---
 benchmark/run.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmark/run.py b/benchmark/run.py
index 2016701..9e7d783 100644
--- a/benchmark/run.py
+++ b/benchmark/run.py
@@ -134,7 +134,7 @@ elif args.format == 'onnx':
         for o_n in output_names:
             dynamic_axes[o_n] = {0: 'batch_size', 1: 'sequence_length'}
         torch.onnx.export(
-            op.model,
+            op.model.model,
             tuple(inputs.values()),
             onnx_path,
             input_names=input_names,