Upload folder using huggingface_hub

modeling_latex_ocr.py

@@ -349,10 +349,10 @@ class CustomDecoder(nn.Module):
             logits   = self._forward_embeds(cur_embeds, cur_mask)  # (B, seq, vocab)
             next_tok = logits[:, -1, :].argmax(dim=-1)             # (B,)
 
+            finished |= (next_tok == eos_id)
             for i in range(batch):
                 if not finished[i]:
                     gen_ids[i].append(next_tok[i].item())
-            finished |= (next_tok == eos_id)
             if finished.all():
                 break