Merge pull request #316 from kylebgorman/slight

Slight simplification in hard attention loop

Author: kylebgorman

yoyodyne/models/base.py

@@ -370,7 +370,6 @@ def predict_step(
                 using beam search, the predictions and scores as a tuple of
                 tensors; if using greedy search, the predictions as a tensor.
         """
-
         if self.beam_width > 1:
             return self(batch)
         else:

yoyodyne/models/hard_attention.py

@@ -216,9 +216,8 @@ def greedy_decode(
         emissions, transitions, state = self.decode_step(
             encoded, mask, symbol, state
         )
-        symbol, likelihood = self._greedy_step(
-            emissions, transitions[:, 0].unsqueeze(1)
-        )
+        likelihood = transitions[:, 0].unsqueeze(1)
+        symbol = self._greedy_step(emissions, likelihood)
         predictions = [symbol]
         # Tracks when each sequence has decoded an END.
         final = torch.zeros(batch_size, device=self.device, dtype=bool)
@@ -233,7 +232,7 @@ def greedy_decode(
             likelihood = likelihood.logsumexp(dim=2, keepdim=True).transpose(
                 1, 2
             )
-            symbol, likelihood = self._greedy_step(emissions, likelihood)
+            symbol = self._greedy_step(emissions, likelihood)
             predictions.append(symbol)
             final = torch.logical_or(final, symbol == special.END_IDX)
             if final.all():
@@ -257,15 +256,12 @@ def _greedy_step(
                 symbol sequence.
 
         Returns:
-            Tuple[torch.Tensor, torch.Tensor]: greedily decoded symbol
-                for current timestep and the current likelihood of the
-                decoded symbol sequence.
+            torch.Tensor: greedily decoded symbol for the current timestep.
         """
         probabilities = likelihood + emissions.transpose(1, 2)
         probabilities = probabilities.logsumexp(dim=2)
         # -> B.
-        symbol = torch.argmax(probabilities, dim=1)
-        return symbol, likelihood
+        return torch.argmax(probabilities, dim=1)
 
     @staticmethod
     def _gather_at_idx(