diff --git a/spacy_llm/cache.py b/spacy_llm/cache.py
index 92f4e80b..3e1559e1 100644
--- a/spacy_llm/cache.py
+++ b/spacy_llm/cache.py
@@ -72,7 +72,7 @@ def initialize(self, vocab: Vocab, task: LLMTask) -> None:
         """
         Initialize cache with data not available at construction time.
         vocab (Vocab): Vocab object.
-        task (ShardingLLMTask): Task.
+        task (LLMTask): Task.
         """
         self._vocab = vocab
         if isinstance(task, PromptTemplateProvider):