refactor: cast Unitary to cp array on GPU

qiboteam · Feb 3, 2025 · dadf2ec · dadf2ec
1 parent 3fac930
commit dadf2ec
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/qibojit/backends/matrices.py b/src/qibojit/backends/matrices.py
@@ -99,9 +99,8 @@ def _cast(self, x, dtype):
 
     # Necessary to avoid https://github.com/qiboteam/qibo/issues/928
     def Unitary(self, u):
-        if isinstance(u, self.cp.ndarray):
-            u = u.get()
-        return super().Unitary(u)
+        dtype = getattr(np, self.dtype)
+        return self._cast(u, dtype=dtype)
 
 
 class CustomCuQuantumMatrices(CustomMatrices):  # pragma: no cover