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(