nielsr commited on
Commit
941b163
1 Parent(s): 8f381d4

Fix code example

Browse files
Files changed (1) hide show
  1. README.md +3 -4
README.md CHANGED
@@ -36,10 +36,9 @@ processor = TrOCRProcessor.from_pretrained('microsoft/trocr-large-stage1')
36
  model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-large-stage1')
37
 
38
  # training
39
- pixel_values = processor(images=image, return_tensors="pt").pixel_values
40
- labels = processor.tokenizer('industrie', return_tensors='pt').input_ids
41
- outputs = model(pixel_values=pixel_values, labels=decoder_input_ids)
42
- loss = outputs.loss
43
  ```
44
 
45
  ### BibTeX entry and citation info
36
  model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-large-stage1')
37
 
38
  # training
39
+ pixel_values = processor(image, return_tensors="pt").pixel_values # Batch size 1
40
+ decoder_input_ids = torch.tensor([[model.config.decoder.decoder_start_token_id]])
41
+ outputs = model(pixel_values=pixel_values, decoder_input_ids=decoder_input_ids)
 
42
  ```
43
 
44
  ### BibTeX entry and citation info