Update optimum/neuron/utils/hub_neuronx_cache.py

Co-authored-by: Michael Benayoun <[email protected]>
huggingface · Jan 26, 2024 · 4765f2f · 4765f2f
1 parent 1b4a731
commit 4765f2f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/neuron/utils/hub_neuronx_cache.py b/optimum/neuron/utils/hub_neuronx_cache.py
@@ -216,7 +216,7 @@ def hub_neuronx_cache(entry: Optional[CacheEntry] = None):
     """A context manager to activate the Hugging Face Hub proxy compiler cache.
 
     Args:
-        entry (`CacheEntry`):
+        entry (`Optional[CacheEntry]`, defaults to `None`):
             An optional dataclass containing metadata associated with the cache session.
             Will create a dedicated entry in the cache registry.
     """