From b9ca3a6a37943b8e8053445d87a7bd7d561b1bef Mon Sep 17 00:00:00 2001
From: Charles Moussa <charles.moussa@pasqal.com>
Date: Mon, 30 Dec 2024 10:39:55 +0100
Subject: [PATCH] [BugFix] circuit.sample fix (#303)

circuit.sample was incorrect as there were two instructions for
obtaining probabilities. We only need one instructions depending on
whether we have a dendity matrix or a state vector.
---
 pyqtorch/circuit.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pyqtorch/circuit.py b/pyqtorch/circuit.py
index 57ea7cf5..41042bbb 100644
--- a/pyqtorch/circuit.py
+++ b/pyqtorch/circuit.py
@@ -100,9 +100,8 @@ def sample(
                     start_dim=0,
                     end_dim=-2,
                 ).t()
-                probs = torch.abs(torch.pow(state, 2))
+                probs = torch.pow(torch.abs(state), 2)
 
-            probs = torch.pow(torch.abs(state), 2)
             if self.readout_noise is not None:
                 probs = self.readout_noise.apply(probs, n_shots)
             counters = list(