premai-io
/

prem-1B-chat

Text Generation

text-generation-inference

Inference Endpoints

Model card Files Files and versions Community

goku commited on May 9

Commit

43b047b

•

1 Parent(s): 0ca6467

update readme

Files changed (1) hide show

README.md +1 -1

README.md CHANGED Viewed

@@ -53,7 +53,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 # Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("premai-io/prem-1B-chat")
-model = AutoModelForCausalLM.from_pretrained('premai-io/prem-1B-chat')
 model = model.to('cuda')
 # Setup terminators

 # Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("premai-io/prem-1B-chat")
+model = AutoModelForCausalLM.from_pretrained('premai-io/prem-1B-chat', torch_dtype=torch.bfloat16)
 model = model.to('cuda')
 # Setup terminators