Fix num_threads flag being ignored

LahiRumesh · Apr 3, 2023 · 25b1487 · 25b1487
1 parent 57440bf
commit 25b1487
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/llama-cli/src/cli_args.rs b/llama-cli/src/cli_args.rs
@@ -126,7 +126,7 @@ pub struct Args {
 }
 impl Args {
     #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
-    pub(crate) fn num_threads(&self) -> usize {
+    pub(crate) fn autodetect_num_threads(&self) -> usize {
         std::process::Command::new("sysctl")
             .arg("-n")
             .arg("hw.perflevel0.physicalcpu")
@@ -137,9 +137,13 @@ impl Args {
     }
 
     #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
-    pub(crate) fn num_threads(&self) -> usize {
+    pub(crate) fn autodetect_num_threads(&self) -> usize {
         num_cpus::get_physical()
     }
+
+    pub(crate) fn num_threads(&self) -> usize {
+        self.num_threads.unwrap_or_else(|| self.autodetect_num_threads())
+    }
 }
 
 fn parse_bias(s: &str) -> Result<TokenBias, String> {