def _use_beamable_mm(self, beam_size):
"""Replace torch.bmm with BeamableMM in attention layers."""
beamable_mm = BeamableMM(beam_size)
for attn in self.attention:
attn.bmm = beamable_mm
def make_generation_fast_(self, beamable_mm_beam_size=None, **kwargs):
"""Replace torch.bmm with BeamableMM."""
if beamable_mm_beam_size is not None:
del self.bmm
self.add_module('bmm', BeamableMM(beamable_mm_beam_size))
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)