From f69e96acf5df8970350754ccf740057bca8f73f1 Mon Sep 17 00:00:00 2001
From: AllentDan <AllentDan@yeah.net>
Date: Tue, 3 Jan 2023 10:04:14 +0800
Subject: [PATCH] fix no test_step error in TorchWrapper

---
 tools/profiler.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tools/profiler.py b/tools/profiler.py
index 822bb83255..81dcc0b996 100644
--- a/tools/profiler.py
+++ b/tools/profiler.py
@@ -80,8 +80,8 @@ def __init__(self, model):
         self.model = model
 
     @TimeCounter.count_time(Backend.PYTORCH.value)
-    def forward(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
+    def test_step(self, *args, **kwargs):
+        return self.model.test_step(*args, **kwargs)
 
 
 def main():