From c119600d6edbe02349b93cf08372409bffd4cf6a Mon Sep 17 00:00:00 2001 From: Harry Stern Date: Mon, 15 Apr 2024 00:50:32 -0400 Subject: [PATCH] Move image tensor to device in trocr example (#2063) Signed-off-by: Harry Stern --- candle-examples/examples/trocr/main.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/candle-examples/examples/trocr/main.rs b/candle-examples/examples/trocr/main.rs index d254e3f0..f857295c 100644 --- a/candle-examples/examples/trocr/main.rs +++ b/candle-examples/examples/trocr/main.rs @@ -115,7 +115,7 @@ pub fn main() -> anyhow::Result<()> { let processor = image_processor::ViTImageProcessor::new(&processor_config); let image = vec![args.image.as_str()]; - let image = processor.preprocess(image)?; + let image = processor.preprocess(image)?.to_device(&device)?; let encoder_xs = model.encoder().forward(&image)?;