small fix for reformer

patrickvonplaten · patrickvonplaten · commit 12b54eceeb57 · 2020-11-03T11:36:23.000Z
diff --git a/src/transformers/generation_utils.py b/src/transformers/generation_utils.py
@@ -189,6 +189,8 @@ def _update_model_kwargs_for_generation(
             model_kwargs["past"] = outputs.past_key_values
         elif "mems" in outputs:
             model_kwargs["past"] = outputs.mems
+        elif "past_buckets_states" in outputs:
+            model_kwargs["past"] = outputs.past_buckets_states
         else:
             model_kwargs["past"] = None
 
