Fixing docs + causal_lm batch_class.

huggingface · Jul 4, 2024 · a1a5ff6 · a1a5ff6
1 parent 26c8d2f
commit a1a5ff6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/docs/source/supported_models.md b/docs/source/supported_models.md
@@ -10,6 +10,7 @@ Text Generation Inference enables serving optimized models on specific hardware
 - [Llama](https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct)
 - [Phi 3](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct)
 - [Gemma](https://huggingface.co/google/gemma-7b)
+- [PaliGemma](https://huggingface.co/google/paligemma-3b-pt-224)
 - [Gemma2](https://huggingface.co/google/gemma2-9b)
 - [Cohere](https://huggingface.co/CohereForAI/c4ai-command-r-plus)
 - [Dbrx](https://huggingface.co/databricks/dbrx-instruct)

diff --git a/server/text_generation_server/models/causal_lm.py b/server/text_generation_server/models/causal_lm.py
@@ -642,7 +642,7 @@ def fallback(
 
     @property
     def batch_type(self) -> Type[CausalLMBatch]:
-        return self.batch
+        return self.batch_class
 
     def forward(
         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None