chore: config_name_to_class uses config.model_type now

huggingface · Apr 10, 2024 · f334bbd · f334bbd
1 parent 2215595
commit f334bbd
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/optimum/tpu/modeling.py b/optimum/tpu/modeling.py
@@ -24,7 +24,8 @@
 from optimum.tpu.modeling_gemma import TpuGemmaForCausalLM
 
 def config_name_to_class(pretrained_model_name_or_path: str):
-    if "gemma" in pretrained_model_name_or_path:
+    config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
+    if config.model_type == "gemma":
         return TpuGemmaForCausalLM
     return BaseAutoModelForCausalLM