From 594c8b038fff216fea9a4effb6d39b871a6812e5 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Tue, 6 Feb 2024 01:59:41 +0000
Subject: [PATCH] cleanup

---
 rllm-cuda/src/llamacpp/tmodel.rs | 23 ++++++++---------------
 rllm-cuda/src/llm/tmodel.rs      |  1 -
 2 files changed, 8 insertions(+), 16 deletions(-)

diff --git a/rllm-cuda/src/llamacpp/tmodel.rs b/rllm-cuda/src/llamacpp/tmodel.rs
index 1927ba5e..3e6efd0d 100644
--- a/rllm-cuda/src/llamacpp/tmodel.rs
+++ b/rllm-cuda/src/llamacpp/tmodel.rs
@@ -208,23 +208,16 @@ impl TModel {
     }
 
     fn sample_argmax(&self, logits: &Tensor) -> u32 {
-        #[cfg(feature = "tch")]
-        {
-            logits.argmax(0, false).int64_value(&[]) as u32
-        }
-        #[cfg(not(feature = "tch"))]
-        {
-            let data = logits.as_slice();
-            let mut top = data[0];
-            let mut top_idx = 0;
-            for (i, x) in data.iter().enumerate() {
-                if *x > top {
-                    top = *x;
-                    top_idx = i;
-                }
+        let data = logits.as_slice();
+        let mut top = data[0];
+        let mut top_idx = 0;
+        for (i, x) in data.iter().enumerate() {
+            if *x > top {
+                top = *x;
+                top_idx = i;
             }
-            top_idx as u32
         }
+        top_idx as u32
     }
 
     fn sample_multinomial(&self, state: &mut LogitsProcessor, prs: &Vec<f32>) -> Result<u32> {
diff --git a/rllm-cuda/src/llm/tmodel.rs b/rllm-cuda/src/llm/tmodel.rs
index dd77276d..987b7cbf 100644
--- a/rllm-cuda/src/llm/tmodel.rs
+++ b/rllm-cuda/src/llm/tmodel.rs
@@ -152,7 +152,6 @@ impl ModelExec for TModel {
         num_seqs: usize,
         vocab_size: usize,
     ) -> Self::AiciBias {
-        #[cfg(feature = "tch")]
         let tensor = Tensor::from_slice(slice)
             .to(self.config.model.device)
             .reshape(&[num_seqs as i64, vocab_size as i64]);