Fix BeamableMM

fairseq/models/fairseq_incremental_decoder.py

@@ -110,7 +110,9 @@ class FairseqIncrementalDecoder(FairseqDecoder):
 
     def set_beam_size(self, beam_size):
         """Sets the beam size in the decoder and all children."""
-        def apply_set_beam_size(module):
-            if module != self and hasattr(module, 'set_beam_size'):
-                module.set_beam_size(beam_size)
-        self.apply(apply_set_beam_size)
+        if getattr(self, '_beam_size', -1) != beam_size:
+            def apply_set_beam_size(module):
+                if module != self and hasattr(module, 'set_beam_size'):
+                    module.set_beam_size(beam_size)
+            self.apply(apply_set_beam_size)
+            self._beam_size = beam_size

fairseq/models/fairseq_model.py

@@ -62,6 +62,11 @@ class FairseqModel(nn.Module):
                 return
         self.apply(apply_remove_weight_norm)
 
+        def apply_make_generation_fast_(module):
+            if module != self and hasattr(module, 'make_generation_fast_'):
+                module.make_generation_fast_(**kwargs)
+        self.apply(apply_make_generation_fast_)
+
         def train(mode):
             if mode:
                 raise RuntimeError('cannot train after make_generation_fast')
@@ -69,8 +74,3 @@ class FairseqModel(nn.Module):
         # this model should no longer be used for training
         self.eval()
         self.train = train
-
-        def apply_make_generation_fast_(module):
-            if module != self and hasattr(module, 'make_generation_fast_'):
-                module.make_generation_fast_(**kwargs)
-        self.apply(apply_make_generation_fast_)

fairseq/models/fconv.py

@@ -145,7 +145,8 @@ class AttentionLayer(nn.Module):
     def make_generation_fast_(self, beamable_mm_beam_size=None, **kwargs):
         """Replace torch.bmm with BeamableMM."""
         if beamable_mm_beam_size is not None:
-            self.bmm = BeamableMM(beamable_mm_beam_size)
+            del self.bmm
+            self.add_module('bmm', BeamableMM(beamable_mm_beam_size))
 
 
 class FConvDecoder(FairseqIncrementalDecoder):