add warm-up method for inference engines

OpenNMT · Apr 4, 2024 · 5f036ab · 5f036ab
1 parent cb1cb22
commit 5f036ab
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/onmt/inference_engine.py b/onmt/inference_engine.py
@@ -155,6 +155,13 @@ def __init__(self, opt):
             self.transforms_cls = get_transforms_cls(opt._all_transform)
             self.vocabs = self.translator.vocabs
 
+    def warm_up(self):
+        from onmt.translate.translator import build_translator
+
+        self.translator = build_translator(
+            self.opt, self.device_id, logger=self.logger, report_score=True
+        )
+
     def _translate(self, infer_iter):
         scores, preds = self.translator._translate(
             infer_iter, infer_iter.transforms, self.opt.attn_debug, self.opt.align_debug