From cc15f3dec1f041cf0ed54447e4d8d5a345b0a8c4 Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincent.michaud-rioux@xanadu.ai>
Date: Tue, 2 Apr 2024 18:23:31 +0000
Subject: [PATCH] Fix setBasisState.

---
 .../src/simulators/lightning_kokkos/StateVectorKokkosMPI.hpp    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkosMPI.hpp b/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkosMPI.hpp
index 66c2bc6c61..dd66c35674 100644
--- a/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkosMPI.hpp
+++ b/pennylane_lightning/core/src/simulators/lightning_kokkos/StateVectorKokkosMPI.hpp
@@ -537,7 +537,7 @@ class StateVectorKokkosMPI final
     void setBasisState(const std::size_t global_index) {
         const auto index = global_2_local_index(global_index);
         const auto rank = static_cast<std::size_t>(get_mpi_rank());
-        if (index.first == 0) {
+        if (index.first == rank) {
             (*sv_).setBasisState(index.second);
         } else {
             (*sv_).initZeros();