Update README.md

README.md

@@ -39,7 +39,7 @@ from transformers import pipeline
 
 generate_text = pipeline(
     model="h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b",
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
     trust_remote_code=True,
     use_fast=False,
     device_map={"": "cuda:0"},
@@ -83,7 +83,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 model = AutoModelForCausalLM.from_pretrained(
     "h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b",
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
     device_map={"": "cuda:0"}
 )
 generate_text = H2OTextGenerationPipeline(model=model, tokenizer=tokenizer)
@@ -113,7 +113,7 @@ model_name = "h2oai/h2ogpt-gm-oasst1-multilang-2048-falcon-7b"  # either local f
 prompt = "<|prompt|>How are you?<|endoftext|><|answer|>"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
 model.cuda().eval()
 inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("cuda")