Update registry.py

ROCm · Dec 2, 2024 · c0e9b46 · c0e9b46
1 parent efb0432
commit c0e9b46
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
@@ -90,6 +90,9 @@
     "BartModel": ("bart", "BartForConditionalGeneration"),
     "BartForConditionalGeneration": ("bart", "BartForConditionalGeneration"),
     "Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"),  # noqa: E501
+    "T5Model": ("t5", "T5ForConditionalGeneration"),
+    "T5ForConditionalGeneration": ("t5", "T5ForConditionalGeneration"),
+    "T5WithLMHeadModel": ("t5", "T5ForConditionalGeneration")
 }
 
 _EMBEDDING_MODELS = {
@@ -484,4 +487,4 @@ def _run() -> None:
 
 
 if __name__ == "__main__":
-    _run()
+    _run()