#0: Minor fix

tenstorrent · Jan 17, 2025 · 1075e20 · 1075e20
1 parent 9bebdc3
commit 1075e20
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/models/demos/llama3/demo/simple_text_demo.py b/models/demos/llama3/demo/simple_text_demo.py
@@ -475,9 +475,7 @@ def test_llama_demo_text(
 
             # Get the next token
             if argmax_on_device:
-                out_tok = logits
-                if out_tok.dim() == 1:
-                    out_tok = out_tok.unsqueeze(0)
+                out_tok = logits.unsqueeze(1)
             else:
                 # TODO Fix use case with temperature > 0
                 _, out_tok = sample_host(