fix(frontend): fix client decryption casting

zama-ai · Oct 15, 2024 · 69e545d · 69e545d
1 parent 056f931
commit 69e545d
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/frontends/concrete-python/concrete/fhe/compilation/client.py b/frontends/concrete-python/concrete/fhe/compilation/client.py
@@ -294,11 +294,10 @@ def decrypt(
         decrypted = tuple(
             client_circuit.process_output(
                 result._inner, position  # pylint: disable=protected-access
-            )
-            .to_py_val()
-            .astype("int64")
+            ).to_py_val()
             for position, result in enumerate(flattened_results)
         )
+        decrypted = tuple(d.astype("int64") if isinstance(d, np.ndarray) else d for d in decrypted)
 
         return decrypted if len(decrypted) != 1 else decrypted[0]
 
@@ -346,11 +345,10 @@ def simulate_decrypt(
         decrypted = tuple(
             client_circuit.simulate_process_output(
                 result._inner, position  # pylint: disable=protected-access
-            )
-            .to_py_val()
-            .astype("int64")
+            ).to_py_val()
             for position, result in enumerate(flattened_results)
         )
+        decrypted = tuple(d.astype("int64") if isinstance(d, np.ndarray) else d for d in decrypted)
 
         return decrypted if len(decrypted) != 1 else decrypted[0]