Clamping beam_idxs, still fails trying to cat past_states to hidden_states

Author: joecummings

torchtext/prototype/generate.py

@@ -298,25 +298,26 @@ def update_func(emissions, N, T, prev_step_token_idxs, prev_step_hyp_idxs, prev_
 
                         # We could store this in model_kwargs
                         num_hyps_in_prev_step = model_kwargs["past"][0][0].shape[0]
-                        
+
                         num_finished_hyps_in_step = num_hyps_in_prev_step - len(prev_step_hyp_idxs)
                         if num_finished_hyps_in_step > 0:
                             beam_idxs = F.pad(beam_idxs, (0, num_finished_hyps_in_step), "constant", 0)
-                        
+
+                        beam_idxs = torch.clamp(beam_idxs, max=len(prev_step_hyp_idxs) - 1)
+
                         reordered_cached = self.model._reorder_cache(model_kwargs["past"], beam_idxs)
 
                         if num_finished_hyps_in_step > 0:
                             sliced_cache = ()
                             for states in reordered_cached:
                                 sliced_state = ()
                                 for state in states:
-                                    sliced_state = sliced_state + (state[:len(prev_step_hyp_idxs)],)
+                                    sliced_state = sliced_state + (state[: len(prev_step_hyp_idxs)],)
                                 sliced_cache = sliced_cache + (sliced_state,)
                             reordered_cached = sliced_cache
 
                         model_inputs["past_key_values"] = reordered_cached
 
-
                 # Forward pass
                 outputs = self.model(**model_inputs)