Hi @nrazavi !
I've already started a similar topic some time ago:
At that time, @srush, @jean.senellart and @Dakun gave me some advice about how to implement it.
Unfortunately, I haven't had time to implement it since the new openNMT version was released including the new BeamSearch
However, I am still interested in this feature too.
I hope it will appear in new releases soon