From 6d6d87f8b3b5be738c253ec9597e2a22150a47af Mon Sep 17 00:00:00 2001
From: laurent <laurent.mazare@gmail.com>
Date: Fri, 19 Apr 2024 14:22:01 +0200
Subject: [PATCH] Use BF16 for llama v3 by default.

---
 candle-examples/examples/llama/main.rs | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/candle-examples/examples/llama/main.rs b/candle-examples/examples/llama/main.rs
index 32763153..b04a1dca 100644
--- a/candle-examples/examples/llama/main.rs
+++ b/candle-examples/examples/llama/main.rs
@@ -120,7 +120,10 @@ fn main() -> Result<()> {
         Some("bf16") => DType::BF16,
         Some("f32") => DType::F32,
         Some(dtype) => bail!("Unsupported dtype {dtype}"),
-        None => DType::F16,
+        None => match args.which {
+            Which::V3 | Which::V3Instruct => DType::BF16,
+            Which::V1 | Which::V2 | Which::Solar10_7B | Which::TinyLlama1_1BChat => DType::F16,
+        },
     };
     let (llama, tokenizer_filename, mut cache, config) = {
         let api = Api::new()?;