Load bfloat16

pomonam · Jun 28, 2024 · d9df4ed · d9df4ed
1 parent b6eaddd
commit d9df4ed
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/openwebtext/pipeline.py b/examples/openwebtext/pipeline.py
@@ -13,7 +13,7 @@
 
 def construct_llama3() -> nn.Module:
     config = AutoConfig.from_pretrained(
-        MODEL_NAME, trust_remote_code=True, device_map="auto"
+        MODEL_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto"
     )
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,