Fix a bug in 1D input shape (#5)

2023-03-06 10:05:27 -08:00
parent 3e9f991d6a
commit 04e5acc08e
3 changed files with 11 additions and 6 deletions
--- a/cacheflow/models/input_metadata.py
+++ b/cacheflow/models/input_metadata.py
@@ -24,7 +24,7 @@ class InputMetadata:

        self.num_prompts = len(prompt_lens)
        self.num_generation_tokens = context_lens.shape[0]
-        self.num_valid_tokens = len(slot_mapping)
+        self.num_valid_tokens = slot_mapping.shape[0]
        if block_tables.numel() > 0:
            self.max_num_blocks_per_seq = block_tables.shape[1]
        else: