From daf3169ef8721e0b0a97ac6a152d4a9b4d43ca0d Mon Sep 17 00:00:00 2001 From: youkaichao Date: Wed, 22 Jan 2025 13:29:23 +0800 Subject: [PATCH] doc string for use_memory_pool Signed-off-by: youkaichao --- vllm/device_allocator/cumem.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index 831936c12460c..374dcb13b598b 100644 --- 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: