Switch to f16

tracel-ai · May 27, 2024 · 59061d1 · 59061d1
1 parent 9492e20
commit 59061d1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llama-burn/examples/generate.rs b/llama-burn/examples/generate.rs
@@ -2,7 +2,7 @@ use std::time::Instant;
 
 use burn::{
     backend::{libtorch::LibTorchDevice, LibTorch},
-    tensor::backend::Backend,
+    tensor::{backend::Backend, f16},
 };
 use clap::Parser;
 use llama_burn::{
@@ -71,7 +71,7 @@ pub fn generate<B: Backend, T: Tokenizer>(
 }
 
 pub fn main() {
-    type B = LibTorch;
+    type B = LibTorch<f16>;
 
     // Parse arguments
     let args = Config::parse();