Move image tensor to device in trocr example (#2063)

Signed-off-by: Harry Stern <harry@harrystern.net>
This commit is contained in:
Harry Stern
2024-04-15 00:50:32 -04:00
committed by GitHub
parent c449f65b12
commit c119600d6e

View File

@ -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)?;