diff --git a/setup.py b/setup.py
index 56c087f85a..0ef706955f 100644
--- a/setup.py
+++ b/setup.py
@@ -158,12 +158,9 @@ def build_extension(self, ext: CMakeExtension):
 
         # Ensure that catalyst shared object is copied to the build directory for pip editable install
         if backend in ("lightning_kokkos"):
-            if platform.system() == "Linux":
-                source = os.path.join(f"{extdir}", f"lib{backend}_catalyst.so")
-                destination = os.path.join(os.getcwd(), "build")
-                shutil.copy(source, destination)
-            if platform.system() == "Darwin":
-                source = os.path.join(f"{extdir}", f"lib{backend}_catalyst.dylib")
+            if platform.system() in ["Linux", "Darwin"]:
+                shared_lib_ext = {"Linux": ".so", "Darwin": ".dylib"}[platform.system()]
+                source = os.path.join(f"{extdir}", f"lib{backend}_catalyst{shared_lib_ext}")
                 destination = os.path.join(os.getcwd(), "build")
                 shutil.copy(source, destination)