diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 3e23ae415f..3ddfb4c608 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -120,7 +120,7 @@ def completion( ) else: hf_model_name = model - if "jumpstart-dft-meta-textgeneration-llama" in model or "meta-textgenerationneuron-llama-2" in model: # llama2 model + if "meta-textgeneration-llama-2" in model or "meta-textgenerationneuron-llama-2" in model: # llama2 model if model.endswith("-f") or "-f-" in model or "chat" in model: # sagemaker default for a chat model hf_model_name = "meta-llama/Llama-2-7b-chat" # apply the prompt template for a llama2 chat model else: