Merge pull request #118 from danbev/with_main_gpu

add with_main_gpu to LlamaModelParams
utilityai · Feb 29, 2024 · 446d16d · 446d16d
2 parents 52687ca + dd264e7
commit 446d16d
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/llama-cpp-2/src/model/params.rs b/llama-cpp-2/src/model/params.rs
@@ -56,6 +56,13 @@ impl LlamaModelParams {
         self
     }
 
+    /// sets the main GPU
+    #[must_use]
+    pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
+        self.params.main_gpu = main_gpu;
+        self
+    }
+
     /// sets `vocab_only`
     #[must_use]
     pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {