diff --git a/qadence/backends/pulser/backend.py b/qadence/backends/pulser/backend.py
index 7eb383de..8d086af3 100644
--- a/qadence/backends/pulser/backend.py
+++ b/qadence/backends/pulser/backend.py
@@ -351,27 +351,14 @@ def expectation(
                     noise=noise,
                 )
                 support = sorted(list(circuit.abstract.register.support))
-                # TODO: There should be a better check for batched density matrices.
-                if dms.size()[0] > 1:
-                    res_list = [
-                        [
-                            obs.native(
-                                dm.squeeze(), param_values, qubit_support=support, noise=noise
-                            )
-                            for dm in dms
-                        ]
-                        for obs in observable
-                    ]
-                    res = torch.stack(
-                        [torch.transpose(torch.stack(res), 0, -1) for res in res_list]
-                    )
-
-                else:
-                    res_list = [
-                        obs.native(dms, param_values, qubit_support=support) for obs in observable
+                res_list = [
+                    [
+                        obs.native(dm.squeeze(), param_values, qubit_support=support, noise=noise)
+                        for dm in dms
                     ]
-                    res = torch.transpose(torch.stack(res_list), 0, 1)
-                res = res if len(res.shape) > 0 else res.reshape(1)
+                    for obs in observable
+                ]
+                res = torch.stack([torch.transpose(torch.stack(res), 0, -1) for res in res_list])
                 return res.real
         elif mitigation is not None:
             logger.warning(
diff --git a/tests/conftest.py b/tests/conftest.py
index fbbb6364..9b0b9f91 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -33,7 +33,7 @@ def batched_noisy_pulser_sim() -> Tensor:
 
 @fixture
 def noisy_pulser_sim() -> Tensor:
-    return tensor([[[0.3043, 0.0000, 0.0000, -0.0038]]])
+    return tensor([[0.3597]])
 
 
 @fixture