From 9cc7436e65e2fb5ef743dfe52e9ad95dee1cdc21 Mon Sep 17 00:00:00 2001 From: asaigal Date: Fri, 1 Nov 2024 21:18:29 -0700 Subject: [PATCH] #0: Modify llama demo to use blocking reads --- models/demos/llama3/demo/demo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/demos/llama3/demo/demo.py b/models/demos/llama3/demo/demo.py index d640e13d707..34b46e35e8b 100644 --- a/models/demos/llama3/demo/demo.py +++ b/models/demos/llama3/demo/demo.py @@ -398,7 +398,7 @@ def run_llama3_demo(user_input, batch_size, mesh_device, instruct_mode, is_ci_en # Write to host ttnn.wait_for_event(1, op_event) tt_output_torch = ttnn.to_torch( - tt_out_tok.cpu(blocking=False, cq_id=1), mesh_composer=ttnn.ConcatMeshToTensor(mesh_device, dim=1) + tt_out_tok.cpu(blocking=True, cq_id=1), mesh_composer=ttnn.ConcatMeshToTensor(mesh_device, dim=1) )[0, 0, 0, :batch_size] ttnn.record_event(1, write_event)