Fix minor bug with token input processing for galaxy

Signed-off-by: Salar <[email protected]>
tenstorrent · Dec 24, 2024 · 27249e2 · 27249e2
1 parent 2fe5001
commit 27249e2
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/models/demos/llama3/tt/llama_model.py b/models/demos/llama3/tt/llama_model.py
@@ -110,9 +110,8 @@ def prepare_inputs_prefill(self, tokens, start_pos=0, page_table=None, chunk_pag
 
         tokens = tokens.reshape(1, 1, 1, -1)
         S = tokens.shape[-1]
-        dims = (None, -1) if self.args.is_galaxy else (None, None)
+        dims = (None, None)  # replicate
         mesh_mapper = ttnn.ShardTensor2dMesh(self.mesh_device, dims=dims, mesh_shape=self.args.cluster_shape)
-
         tokens = ttnn.from_torch(
             tokens,
             device=self.mesh_device,