Update README.md

README.md

@@ -64,7 +64,7 @@ import torch
 
 model_name_or_id = "MaralGPT/Maral-7B-alpha-1"
 
-model = AutoModelForCausalLM.from_pretrained(model_name_or_id, torch_dtype=torch.float16, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_name_or_id, torch_dtype=torch.bfloat16, device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_id)
 
 prompt = "در سال ۱۹۹۶ چه کسی رییس جمهور آمریکا بود؟"
@@ -89,7 +89,7 @@ print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 The code is pretty much the same as above, but with a slight diferrence. 
 
 * Make sure `bitsandbytes` is installed correctly.
-* Your model loading must be `model = AutoModelForCausalLM.from_pretrained(model_name_or_id, load_in_8bit=True, torch_dtype=torch.float16, device_map="auto")`
+* Your model loading must be `model = AutoModelForCausalLM.from_pretrained(model_name_or_id, load_in_8bit=True, torch_dtype=torch.bfloat16, device_map="auto")`
 
 On _free version_ of Google Colab, you may face RAM problems. I guess using `low_cpu_mem_usage=True` in model loading would help.