diff options
Diffstat (limited to 'candle-examples/examples/stable-lm/main.rs')
| -rw-r--r-- | candle-examples/examples/stable-lm/main.rs | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/candle-examples/examples/stable-lm/main.rs b/candle-examples/examples/stable-lm/main.rs index 0535aa70..ccd924a4 100644 --- a/candle-examples/examples/stable-lm/main.rs +++ b/candle-examples/examples/stable-lm/main.rs @@ -234,13 +234,14 @@ fn main() -> Result<()> { let start = std::time::Instant::now(); let config = Config::stablelm_3b_4e1t(args.use_flash_attn); + let device = candle_examples::device(args.cpu)?; let (model, device) = if args.quantized { let filename = &filenames[0]; - let vb = candle_transformers::quantized_var_builder::VarBuilder::from_gguf(filename)?; + let vb = + candle_transformers::quantized_var_builder::VarBuilder::from_gguf(filename, &device)?; let model = QStableLM::new(&config, vb)?; (Model::Quantized(model), Device::Cpu) } else { - let device = candle_examples::device(args.cpu)?; let dtype = if device.is_cuda() { DType::BF16 } else { |