doc string for use_memory_pool

Signed-off-by: youkaichao <[email protected]>
vllm-project · Jan 22, 2025 · daf3169 · daf3169
1 parent d1991e5
commit daf3169
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py
@@ -217,7 +217,12 @@ def use_memory_pool(self, tag: Optional[str] = None):
         """
         A context manager to use the memory pool.
         All memory allocation created inside the context will be allocated 
-        in the memory pool, and has the specified tag."""
+        in the memory pool, and has the specified tag.
+
+        Args:
+            tag: The tag of the memory allocation. If None, the default tag 
+            will be used.
+        """
         if tag is None:
             tag = CuMemAllocator.default_tag
         else: