Support refinement iterations on ONNX (Fixes baudm#12 baudm#66)

Author: baudm

README.md

@@ -91,7 +91,7 @@ print('Decoded label = {}'.format(label[0]))
 
 ## Frequently Asked Questions
 - How do I train on a new language? See Issues [#5](https://github.com/baudm/parseq/issues/5) and [#9](https://github.com/baudm/parseq/issues/9).
-- Can you export to TorchScript or ONNX? Yes, with a [caveat (Issue #12)](https://github.com/baudm/parseq/issues/12#issuecomment-1267842315).
+- Can you export to TorchScript or ONNX? Yes, see Issue [#12](https://github.com/baudm/parseq/issues/12#issuecomment-1267842315).
 - How do I test on my own dataset? See Issue [#27](https://github.com/baudm/parseq/issues/27).
 - How do I finetune and/or create a custom dataset? See Issue [#7](https://github.com/baudm/parseq/issues/7).
 - What is `val_NED`? See Issue [#10](https://github.com/baudm/parseq/issues/10).

strhub/models/parseq/system.py

@@ -144,7 +144,7 @@ def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
             for i in range(self.refine_iters):
                 # Prior context is the previous output.
                 tgt_in = torch.cat([bos, logits[:, :-1].argmax(-1)], dim=1)
-                tgt_padding_mask = ((tgt_in == self.eos_id).cumsum(-1) > 0)  # mask tokens beyond the first EOS token.
+                tgt_padding_mask = ((tgt_in == self.eos_id).int().cumsum(-1) > 0)  # mask tokens beyond the first EOS token.
                 tgt_out = self.decode(tgt_in, memory, tgt_mask, tgt_padding_mask,
                                       tgt_query=pos_queries, tgt_query_mask=query_mask[:, :tgt_in.shape[1]])
                 logits = self.head(tgt_out)