make autograd test compatible with both

PennyLaneAI · Apr 12, 2024 · 9301d22 · 9301d22
1 parent 17ed5ff
commit 9301d22
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
@@ -200,8 +200,12 @@ def test_no_batch_transform(self, mocker):
         tape = qml.tape.QuantumScript.from_queue(q)
         spy = mocker.spy(dev, "batch_transform")
 
-        res = qml.execute([tape], dev, None, device_batch_transform=False)
-        assert np.allclose(res[0], np.cos(y), atol=0.1)
+        if qml.operation.active_new_opmath():
+            res = qml.execute([tape], dev, None, device_batch_transform=False)
+            assert np.allclose(res[0], np.cos(y), atol=0.1)
+        else:
+            with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"):
+                qml.execute([tape], dev, None, device_batch_transform=False)
 
         spy.assert_not_called()