microsoft

/

trocr-base-printed

like 149

Follow

Microsoft 5,104

Hi,

I am trying to trace this model like:

traced_script_module = torch.jit.trace(model, example_input)
traced_script_module_optimized = optimize_for_mobile(traced_script_module)
traced_script_module_optimized._save_for_lite_interpreter("model.ptl")

However, the model consists of a split preprocess and decoder stage (from the docs):

processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-printed')
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-printed')
pixel_values = processor(images=image, return_tensors="pt").pixel_values

generated_ids = model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

How can I trace this model?