Make Captum log_usage universal by supporting async functions

Summary: As title. Checks if function being wrapped is a coroutine and uses an async wrapper instead of sync when logging.

Differential Revision: D76837949

Author: craymichael

captum/log/dummy_log.py

@@ -2,7 +2,6 @@
 
 # pyre-strict
 
-from functools import wraps
 from types import TracebackType
 from typing import Any, List, Optional, Union
 
@@ -32,13 +31,7 @@ def log_usage(*log_args: Any, **log_kwargs: Any):
     # pyre-fixme[3]: Return type must be annotated.
     # pyre-fixme[2]: Parameter must be annotated.
     def _log_usage(func):
-        @wraps(func)
-        # pyre-fixme[53]: Captured variable `func` is not annotated.
-        # pyre-fixme[3]: Return type must be annotated.
-        def wrapper(*args: Any, **kwargs: Any):
-            return func(*args, **kwargs)
-
-        return wrapper
+        return func
 
     return _log_usage