Add padding="longest"

main.py

@@ -164,6 +164,9 @@ def GetChatModel(compute_mode, ckpt_dir=None):
             return_full_text=False,
             # It seems that max_new_tokens has to be specified here, not in .invoke()
             max_new_tokens=2000,
+            # Use padding for FlashAttention alignment
+            # https://github.com/google-deepmind/gemma/issues/169
+            padding="longest",
         )
         # We need the task so HuggingFacePipeline can deal with our class
         pipe.task = "text-generation"

pipeline.py

@@ -22,7 +22,6 @@ class MyTextGenerationPipeline(TextGenerationPipeline):
         continue_final_message=None,
         **generate_kwargs,
     ):
-        print(f"PADDING: {padding}")
         # Only set non-None tokenizer kwargs, so as to rely on the tokenizer's defaults
         tokenizer_kwargs = {
             "add_special_tokens": add_special_tokens,