Fix run_decoders bug (#11871)

intel-analytics · Aug 20, 2024 · bdaeee1 · bdaeee1
1 parent 32f0a77
commit bdaeee1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/llm/src/ipex_llm/transformers/npu_models/llama_mp.py b/python/llm/src/ipex_llm/transformers/npu_models/llama_mp.py
@@ -601,7 +601,7 @@ def forward(
                 x_np[2].ctypes.data_as(ctypes.c_void_p),
             )
             t0 = time.perf_counter()
-            backend_lib.run_decoders(models_ptr, inputs_ptr, 2, 3)
+            backend_lib.run_decoders(models_ptr, inputs_ptr, self.intra_stages, 3)
             t1 = time.perf_counter()
 
         hidden_states = self.backend_decoders[-1].torch_out[0]