diff --git a/demonstrations/ahs_aquila.py b/demonstrations/ahs_aquila.py
index c579b2f385..17fea98493 100644
--- a/demonstrations/ahs_aquila.py
+++ b/demonstrations/ahs_aquila.py
@@ -503,6 +503,8 @@ def gaussian_fn(p, t):
 # the gaussian pulse envelope via trial-and-error to find settings that result in a pi-pulse:
 #
 
+from functools import partial
+
 max_amplitude = 2.0
 displacement = 1.0
 sigma = 0.3
@@ -512,9 +514,10 @@ def gaussian_fn(p, t):
 params = [amplitude_params]
 ts = [0.0, 1.75]
 
-default_qubit = qml.device("default.qubit", wires=3, shots=1000)
+default_qubit = qml.device("default.qubit", wires=3)
 
 
+@partial(qml.set_shots, shots=1000)
 @qml.qnode(default_qubit, interface="jax")
 def circuit(parameters):
     qml.evolve(global_drive)(parameters, ts)
diff --git a/demonstrations/covalent_cloud_gpu.py b/demonstrations/covalent_cloud_gpu.py
index 2dbcf02894..439ee11edc 100644
--- a/demonstrations/covalent_cloud_gpu.py
+++ b/demonstrations/covalent_cloud_gpu.py
@@ -33,6 +33,7 @@
 import covalent_cloud as cc
 import matplotlib.pyplot as plt
 import pennylane as qml
+from functools import partial
 from matplotlib.colors import ListedColormap
 from pennylane import numpy as np
 from sklearn.datasets import make_blobs
@@ -124,7 +125,8 @@
 
 def get_kernel_circuit(n_wires):
 
-    @qml.qnode(qml.device(QML_DEVICE, wires=n_wires, shots=None))
+    @partial(qml.set_shots, shots=None)
+    @qml.qnode(qml.device(QML_DEVICE, wires=n_wires))
     def circuit(x1, x2):
         qml.IQPEmbedding(x1, wires=range(n_wires), n_repeats=4)
         qml.adjoint(qml.IQPEmbedding)(x2, wires=range(n_wires), n_repeats=4)