add .to_dtype for input tensor dtype conversion

atoma-network · Apr 10, 2024 · 13f2d2b · 13f2d2b
1 parent 385c757
commit 13f2d2b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/atoma-inference/src/jrpc_server/mod.rs b/atoma-inference/src/jrpc_server/mod.rs
@@ -1,6 +1,6 @@
-use std::{net::Shutdown, sync::Arc};
+use std::sync::Arc;
 
-use axum::{extract::State, http::StatusCode, routing::post, Extension, Json, Router};
+use axum::{http::StatusCode, routing::post, Extension, Json, Router};
 use serde_json::{json, Value};
 use tokio::sync::{mpsc, oneshot};
 

diff --git a/atoma-inference/src/models/candle/mamba.rs b/atoma-inference/src/models/candle/mamba.rs
@@ -178,7 +178,7 @@ impl ModelTrait for MambaModel {
         let mut output = String::new();
 
         for &token in tokens.iter() {
-            let input = Tensor::new(&[token], &self.device)?;
+            let input = Tensor::new(&[token], &self.device)?.to_dtype(self.dtype)?;
             let logits = self.model.forward(&input, &mut state)?;
 
             next_logits = Some(logits);
@@ -213,7 +213,7 @@ impl ModelTrait for MambaModel {
                 output.push_str(t.as_str());
             }
 
-            let input = Tensor::new(&[next_token], &self.device)?;
+            let input = Tensor::new(&[next_token], &self.device)?.to_dtype(self.dtype)?;
             next_logits = Some(self.model.forward(&input, &mut state)?);
         }
         let dt = start_gen.elapsed();