Skip to content

Commit 765bbe2

Browse files
authored
fix(examples/reranker): add the expected disable_gpu arg for CUDA and Vulkan
1 parent 96b6bcc commit 765bbe2

File tree

1 file changed

+8
-1
lines changed

1 file changed

+8
-1
lines changed

examples/reranker/src/main.rs

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
//! This is a translation of embedding.cpp in llama.cpp using llama-cpp-2.
1+
//! This is an example of reranking documents for a query using llama-cpp-2.
22
#![allow(
33
clippy::cast_possible_wrap,
44
clippy::cast_possible_truncation,
@@ -45,6 +45,11 @@ struct Args {
4545
/// Whether to normalise the produced embeddings
4646
#[clap(long, default_value_t = true)]
4747
normalise: bool,
48+
49+
/// Disable offloading layers to the gpu
50+
#[cfg(any(feature = "cuda", feature = "vulkan"))]
51+
#[clap(long)]
52+
disable_gpu: bool,
4853
}
4954

5055
fn main() -> Result<()> {
@@ -54,6 +59,8 @@ fn main() -> Result<()> {
5459
documents,
5560
pooling,
5661
normalise,
62+
#[cfg(any(feature = "cuda", feature = "vulkan"))]
63+
disable_gpu,
5764
} = Args::parse();
5865

5966
// init LLM

0 commit comments

Comments
 (0)