Fix reset prefix caching

toslali-ibm · toslali-ibm · commit 20d6fefef96b · 2025-03-24T11:32:53.000-04:00
diff --git a/trl/extras/vllm_coloc_client.py b/trl/extras/vllm_coloc_client.py
@@ -129,5 +129,5 @@ def reset_prefix_cache(self):
         """
         Resets the prefix cache for the model.
         """
-        self.llm.reset()
+        self.llm.reset_prefix_cache()
 
