How to properly trace this model?
#4
by z0rz0rz0r - opened
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?