update batch params

PennyLaneAI · Nov 9, 2021 · 859d9d7 · 859d9d7
1 parent 5a4d8d8
commit 859d9d7
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
@@ -41,7 +41,7 @@ def circuit(data, x, weights):
 
     spy = mocker.spy(circuit.device, "batch_execute")
     res = circuit(data, x, weights)
-    assert res.shape == (batch_size, 1, 4)
+    assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
 
 
@@ -61,7 +61,7 @@ def circuit(data):
 
     spy = mocker.spy(circuit.device, "batch_execute")
     res = circuit(data)
-    assert res.shape == (batch_size, 1, 4)
+    assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
 
 
@@ -244,7 +244,7 @@ def circuit(x, weights):
 
     spy = mocker.spy(circuit.device, "batch_execute")
     res = circuit(x, weights)
-    assert res.shape == (batch_size, 1, 4)
+    assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size