From 1134a7c732f76ba1c44af57e8d5834f4bf3dcfb9 Mon Sep 17 00:00:00 2001
From: drisspg <drisspguessous@gmail.com>
Date: Mon, 3 Jun 2024 11:16:07 -0700
Subject: [PATCH] add doctsring to profiler

---
 transformer_nuggets/utils/benchmark.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/transformer_nuggets/utils/benchmark.py b/transformer_nuggets/utils/benchmark.py
index 27029d7..a823245 100644
--- a/transformer_nuggets/utils/benchmark.py
+++ b/transformer_nuggets/utils/benchmark.py
@@ -177,7 +177,18 @@ def profiler(
     profile_memory: bool = False,
     with_stack: bool = True,
 ):
-    """Thin wrapper around torch.profiler"""
+    """Thin wrapper around torch.profiler
+
+    Args:
+        path: The path to save the trace file to
+        record_shapes: Record shapes of tensors
+        profile_memory: Profile memory usage
+        with_stack: Record stack traces
+
+    Usage:
+        with profiler(Path("trace.json")):
+            # code to profile
+    """
     assert path.is_file(), f"{path} is not a file."
     path = path.with_suffix(".json")
     # make parent dir if it doesn't exist