Merge pull request #40 from cirKITers/barriers

Barriers
cirKITers · Sep 21, 2024 · f46eb49 · f46eb49
2 parents 2d4dc55 + 12abae1
commit f46eb49
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/qml_essentials/model.py b/qml_essentials/model.py
@@ -326,6 +326,8 @@ def _circuit(
                         wires=q,
                     )
 
+            qml.Barrier(wires=list(range(self.n_qubits)), only_visual=True)
+
         if self.data_reupload:
             self.pqc(params[-1], self.n_qubits)
 
@@ -367,10 +369,11 @@ def _draw(self, inputs=None, figure=False) -> None:
         if isinstance(self.circuit, qml.qnn.torch.TorchLayer):
             # TODO: throws strange argument error if not catched
             return ""
-        if not figure:
-            result = qml.draw(self.circuit)(params=self.params, inputs=inputs)
-        else:
+
+        if figure:
             result = qml.draw_mpl(self.circuit)(params=self.params, inputs=inputs)
+        else:
+            result = qml.draw(self.circuit)(params=self.params, inputs=inputs)
         return result
 
     def draw(self, inputs=None, figure=False) -> None:

diff --git a/tests/test_model.py b/tests/test_model.py
@@ -477,7 +477,8 @@ def test_local_and_global_meas() -> None:
 
         assert (
             out.shape == test_case["out_shape"]
-        ), f"Expected {test_case['out_shape']}, got shape {out.shape} for test case {test_case}"
+        ), f"Expected {test_case['out_shape']}, got shape {out.shape}\
+            for test case {test_case}"
 
 
 def test_parity() -> None: