Stop prediction when EOS is reached

Hi Dongjun,
In line 38, of the Graph class, the following loop continues until max sequence length is decoded.             
            for i in range(2, Config.data.max_seq_length):

Is it possible to break the loop when EOS(end of the sequence) is reached for the batch?

Thank you,
Surendra

