Quick fix when language_id in in_map for pipelined model.

yzhang87 · copybara-github · commit 8ab5f1a16edc · 2022-03-20T20:28:54.000-07:00
PiperOrigin-RevId: 436102423
diff --git a/lingvo/core/conformer_layer.py b/lingvo/core/conformer_layer.py
@@ -965,6 +965,8 @@ def _FProp(self, theta, in_nmap):
       # TODO(ngyuzh,yuanzx): fix pipelined layer so don't need pass it manually.
       if 'language_vector' in in_nmap:
         out_nmap.language_vector = in_nmap.language_vector
+      if 'language_id' in in_nmap:
+        out_nmap.language_id = in_nmap.language_id
       inputs = self.final_ln.FProp(theta.final_ln, inputs)
       inputs, paddings = self._CastToFPropDtype((inputs, paddings))
       out_nmap.features = inputs
