fix calc fct for inference time

utils.py

@@ -6,18 +6,22 @@ Author:
     - @ChainYo - https://github.com/ChainYo
 """
 
+from contextlib import contextmanager
 from time import time
+from typing import List
 
 
-def timeit(func):
-    """Timeit decorator."""
+@contextmanager
+def calculate_inference_time(buffer: List[int]):
+    """Calculate inference time.
 
-    def wrapper(*args, **kwargs):
-        """Wrapper."""
-        start = time()
-        result = func(*args, **kwargs)
-        end = time()
-        print(f"{func.__name__} took {end - start:.2f} seconds.")
-        return result
+    Args:
+        buffer (list): List of inference times.
 
-    return wrapper
+    Returns:
+        float: Average inference time.
+    """
+    start = time()
+    yield
+    end = time()
+    buffer.append(end - start)