diff --git a/atoma-inference/src/models/candle/falcon.rs b/atoma-inference/src/models/candle/falcon.rs
index 3ef1aa68..59fcbc8b 100644
--- a/atoma-inference/src/models/candle/falcon.rs
+++ b/atoma-inference/src/models/candle/falcon.rs
@@ -273,8 +273,8 @@ mod tests {
             repeat_penalty,
             repeat_last_n,
             max_tokens,
-            top_k,
-            top_p,
+            Some(top_k),
+            Some(top_p),
         );
         let output = model.run(input).expect("Failed to run inference");
 
@@ -355,8 +355,8 @@ mod tests {
             repeat_penalty,
             repeat_last_n,
             max_tokens,
-            top_k,
-            top_p,
+            Some(top_k),
+            Some(top_p),
         );
         let output = model.run(input).expect("Failed to run inference");
         println!("{output}");
diff --git a/atoma-inference/src/models/candle/mistral.rs b/atoma-inference/src/models/candle/mistral.rs
index bf4c907c..0fd66ca7 100644
--- a/atoma-inference/src/models/candle/mistral.rs
+++ b/atoma-inference/src/models/candle/mistral.rs
@@ -263,8 +263,8 @@ mod tests {
             repeat_penalty,
             repeat_last_n,
             max_tokens,
-            top_k,
-            top_p,
+            Some(top_k),
+            Some(top_p),
         );
         let output = model.run(input).expect("Failed to run inference");
 
@@ -351,8 +351,8 @@ mod tests {
             repeat_penalty,
             repeat_last_n,
             max_tokens,
-            top_k,
-            top_p,
+            Some(top_k),
+            Some(top_p),
         );
         let output = model.run(input).expect("Failed to run inference");