Use image_seq_len to determine proper cache size in generate_v2 (#1674)

pytorch · Sep 25, 2024 · f6973d7 · f6973d7
1 parent 2c18c5b
commit f6973d7
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/recipes/dev/generate_v2.py b/recipes/dev/generate_v2.py
@@ -133,9 +133,7 @@ def generate(self, cfg: DictConfig):
                 batch_size=1,
                 dtype=self._dtype,
                 encoder_max_seq_len=(
-                    model_inputs["encoder_mask"][0].size(1)
-                    if is_multimodal_input
-                    else None
+                    self.model_transform.image_seq_len if is_multimodal_input else None
                 ),
                 decoder_max_seq_len=total_response_length,
             )