Add profiling for backward (#120540)

Summary: Fixes #ISSUE_NUMBER X-link: pytorch/pytorch#120540 Approved by: https://github.com/anijain2305 Reviewed By: kit1980 Differential Revision: D54192278 Pulled By: mlazos fbshipit-source-id: 7e813e112f5ce7b36c1fe62f2018494d495b5dc2
pytorch · Feb 27, 2024 · a08683b · a08683b
1 parent 3c38692
commit a08683b
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py
@@ -127,6 +127,12 @@ def tabulate(rows, headers):
         )
 
 
+def maybe_cprofile(func):
+    if config.cprofile:
+        return cprofile_wrapper(func)
+    return func
+
+
 def cprofile_wrapper(func):
     @wraps(func)
     def profile_wrapper(*args, **kwargs):