Hello,
I am working on building a nn.Module class that has a circuit, and that uses qml.qnn.TorchLayer.
For training, I would like to pass batches of images through this layer. This is my code:
class MyNet(nn.Module):
def __init__(self, n_qubits, n_layers=1, device='lightning.qubit', rot_axis='X):
super().__init__()
self.n_qubits = n_qubits
self.rot_axis = rot_axis
self.device = device
self.dev = qml.device(self.device, wires=self.n_qubits)
self.qnode = qml.QNode(self.circuit, self.dev, interface='torch')
self.weight_shapes = {"weights": qml.StronglyEntanglingLayers.shape(n_layers=n_layers, n_wires=self.n_qubits)}
self.qlayer = qml.qnn.TorchLayer(self.qnode, self.weight_shapes)
@qml.batch_params
def circuit(self, inputs, weights):
qml.AngleEmbedding(features=inputs, wires=range(self.n_qubits), rotation=self.rot_axis)
qml.StronglyEntanglingLayers(weights=weights, wires=range(self.n_qubits), imprimitive=qml.ops.CZ)
return [qml.expval(qml.PauliZ(i)) for i in range(self.n_qubits)]
def forward(self, x):
return self.qlayer(x)
I tried using the @qml.batch_params decorator in my code, but it gives me the error:
TypeError: QNode must include an argument with name inputs for inputting data
How can I make passing batches of data work through this self.qlayer?
Why is the decorator not working?
Thank you.