return state for both noisy and non noisy

pyqtorch/hamiltonians/evolution.py

@@ -440,9 +440,8 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
             )
 
             # Retrieve the last state of shape (2**n_qubits, batch_size)
-            state = sol.states[-1]
-
-            return state.reshape([2] * n_qubits + [batch_size])
+            # and reshape
+            state = sol.states[-1].reshape([2] * n_qubits + [batch_size])
         else:
             if not isinstance(state, DensityMatrix):
                 state = density_mat(state)
@@ -454,8 +453,9 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 self.solver_type,
                 options={"use_sparse": self.use_sparse},
             )
-            mesolve_state = sol.states[-1]
-            return mesolve_state
+            # Retrieve the last density matrix
+            state = sol.states[-1]
+        return state
 
     def forward(
         self,