add torch_dtype="auto" to load model weights in bf16

README.md

@@ -41,7 +41,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("cerebras/btlm-3b-8k-base")
-model = AutoModelForCausalLM.from_pretrained("cerebras/-3b-8k-base", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("cerebras/-3b-8k-base", trust_remote_code=True, torch_dtype="auto")
 
 # Set the prompt for generating text
 prompt = "Albert Einstein was known for "
@@ -72,7 +72,7 @@ from transformers import pipeline
 
 # Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("cerebras/btlm-3b-8k-base")
-model = AutoModelForCausalLM.from_pretrained("cerebras/btlm-3b-8k-base", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("cerebras/btlm-3b-8k-base", trust_remote_code=True, torch_dtype="auto")
 
 # Set the prompt for text generation
 prompt = """Isaac Newton was a """