Hello @imakash !
As I mentioned in this post reply, I would suggest first updating your PennyLane installed version.
I also did some minor modifications to the code you shared:
import pennylane as qml
from pennylane import numpy as np
import time as time
import torch
import torch.nn as nn
dev = qml.device("default.qubit", wires=1)
@qml.qnode(dev, interface = 'torch')
def simple_qubit_circuit(inputs, theta):
qml.RX(inputs, wires=0)
qml.RY(theta, wires=0)
return qml.expval(qml.PauliZ(0))
class QNet(nn.Module):
def __init__(self):
super().__init__()
shapes = {
"theta": ()
}
self.q = qml.qnn.TorchLayer(simple_qubit_circuit, shapes)
def forward(self, input_value):
return self.q(input_value)
x_train = np.array([0.2, 0.1, 0.2, 0.14, 0.11, 0.41, 0.55, 0.3, 0.31, 0.6])
x_train = torch.tensor(x_train)
model = QNet()
t1 = time.time()
out = model(x_train)
print("time taken for batch operations: {} seconds".format(time.time()-t1))
and I got:
time taken for batch operations: 0.005563020706176758 seconds
The time seems reasonable to me.
Does it help?