Fix loss casting

qiboteam · Jun 10, 2021 · 2088340 · 2088340
1 parent bc0fce5
commit 2088340
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/qibo/models/variational.py b/src/qibo/models/variational.py
@@ -78,9 +78,11 @@ def _loss(params, circuit, hamiltonian):
             for gate in self.circuit.queue:
                 _ = gate.cache
             loss = K.compile(_loss)
-        elif method != "sgd":
+        elif method == "cma":
             dtype = getattr(K.np, K._dtypes.get("DTYPE"))
             loss = lambda p, c, h: dtype(K.to_numpy(_loss(p, c, h)))
+        elif method != "sgd":
+            loss = lambda p, c, h: K.to_numpy(_loss(p, c, h))
         else:
             loss = _loss