Map state on coverage if coverage is there. (OpenNMT#1338)

AMusketeer · Mar 5, 2019 · 5f08809 · 5f08809
1 parent 90023eb
commit 5f08809
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/onmt/decoders/decoder.py b/onmt/decoders/decoder.py
@@ -182,6 +182,8 @@ def _fix_enc_hidden(hidden):
     def map_state(self, fn):
         self.state["hidden"] = tuple(fn(h, 1) for h in self.state["hidden"])
         self.state["input_feed"] = fn(self.state["input_feed"], 1)
+        if self._coverage and self.state["coverage"] is not None:
+            self.state["coverage"] = fn(self.state["coverage"], 1)
 
     def detach_state(self):
         self.state["hidden"] = tuple(h.detach() for h in self.state["hidden"])