test batchsize and reshape necessary

pasqal-io · Nov 15, 2024 · 7f4adb3 · 7f4adb3
1 parent 64fa4ed
commit 7f4adb3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/pyqtorch/hamiltonians/evolution.py b/pyqtorch/hamiltonians/evolution.py
@@ -454,7 +454,7 @@ def Ht(t: torch.Tensor) -> torch.Tensor:
                 options={"use_sparse": self.use_sparse},
             )
             # Retrieve the last density matrix
-            state = sol.states[-1]
+            state = sol.states[-1].reshape(state.shape)
         return state
 
     def forward(

diff --git a/tests/test_analog.py b/tests/test_analog.py
@@ -353,6 +353,7 @@ def test_timedependent(
     "batchsize",
     [
         1,
+        3,
     ],
 )
 def test_timedependent_with_noise(
-Original file line number
+Diff line change
@@ Expand Up / @@ -353,6 +353,7 @@ def test_timedependent( @@
         "batchsize",
         [
 ,
+,
         ],
     )
     def test_timedependent_with_noise(
@@ Expand Down @@