In ShardedLlamaTest cange seq_lens type to torch.int64

nod-ai · Oct 2, 2024 · 5375807 · 5375807
1 parent 0605ddf
commit 5375807
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/tests/models/llama/sharded_llama_test.py b/sharktank/tests/models/llama/sharded_llama_test.py
@@ -58,7 +58,7 @@ def setUp(self):
             vocab_size=self.vocabulary_size,
         )
         self.prefill_seq_lens = torch.tensor(
-            [14, 9, self.block_seq_stride - 1], dtype=torch.int32
+            [14, 9, self.block_seq_stride - 1], dtype=torch.int64
         )
 
     def make_prefill_args(self, model: PagedLlamaModelV1) -> Dict[str, Any]: