Bug fix

python/fedml/utils/decorators.py

@@ -38,15 +38,17 @@ def wrapper(*args, **kwargs):
 
     return decorator(_func)
 
+
 def timeit(func: Callable):
     """Print the runtime of the decorated function"""
     functools.wraps(func)
 
     def wrapper(*args, **kwargs):
         start = time.perf_counter()
-        func(*args, **kwargs)
+        value = func(*args, **kwargs)
         end = time.perf_counter()
         run_time = end - start
         print(f"Finished {func.__name__!r} in {run_time:.4f} seconds")
+        return value
 
     return wrapper