address comment

This commit is contained in:
Sourab Mangrulkar 2024-03-02 00:43:33 +05:30
parent 15f233b9a1
commit ee5b171250

View file

@ -222,6 +222,8 @@ class Model:
return NomicBertModel return NomicBertModel
if model_architecture == "GemmaForCausalLM": if model_architecture == "GemmaForCausalLM":
return GemmaModel return GemmaModel
if model_architecture == "Starcoder2ForCausalLM":
return Model
return Model return Model
def _is_model_safetensors(self) -> bool: def _is_model_safetensors(self) -> bool: