Problems with gradient computing on pennylane lightning

Hey!
I am trying to train a hybrid torch model , it runs fine with default.qubit but with lightning.qubit or lightning.gpu or even qiskit.remote it fails , I observed two errors , and both of them originates from using qml.probs instead of qml.expval , the First error would be NotImplementedError: Computing the gradient of broadcasted tapes with respect to the broadcasted parameters using the parameter-shift rule gradient transform is currently not supported. See #4462 for details.
the code:

import pennylane as qml
import torch
import numpy as np
from sklearn.datasets import make_moons

X, y = make_moons(n_samples=200, noise=0.1)
y_ = torch.unsqueeze(torch.tensor(y), 1) # used for one-hot encoded labels
y_hot = torch.scatter(torch.zeros((200, 2)), 1, y_, 1)

n_qubits = 2
dev = qml.device(“lightning.qubit”, wires=n_qubits)

@qml.qnode(dev)
def qnode(inputs, weights):
qml.AngleEmbedding(inputs, wires=range(n_qubits))
qml.BasicEntanglerLayers(weights, wires=range(n_qubits))
return qml.probs(wires=[wire for wire in range(n_qubits)])

n_layers = 6
weight_shapes = {“weights”: (n_layers, n_qubits)}

qlayer = qml.qnn.TorchLayer(qnode, weight_shapes)

clayer_1 = torch.nn.Linear(2, 2)
clayer_2 = torch.nn.Linear(4, 2)
softmax = torch.nn.Softmax(dim=1)
layers = [clayer_1, qlayer, clayer_2, softmax]
model = torch.nn.Sequential(*layers)

opt = torch.optim.SGD(model.parameters(), lr=0.2)
loss = torch.nn.L1Loss()

X = torch.tensor(X, requires_grad=True).float()
y_hot = y_hot.float()

batch_size = 5
batches = 200 // batch_size

data_loader = torch.utils.data.DataLoader(
list(zip(X, y_hot)), batch_size=5, shuffle=True, drop_last=True
)

epochs = 6

for epoch in range(epochs):

running_loss = 0

for xs, ys in data_loader:
    opt.zero_grad()

    loss_evaluated = loss(model(xs) , ys)
    loss_evaluated.backward()

    opt.step()

    running_loss += loss_evaluated

avg_loss = running_loss / batches
print("Average loss over epoch {}: {:.4f}".format(epoch + 1, avg_loss))

y_pred = model(X)
predictions = torch.argmax(y_pred, axis=1).detach().numpy()

correct = [1 if p == p_true else 0 for p, p_true in zip(predictions, y)]
accuracy = sum(correct) / len(correct)
print(f"Accuracy: {accuracy * 100}%") 

the second error:

RuntimeError: Function ExecuteTapesBackward returned an invalid gradient at index 0 - got but expected shape compatible with [5]

it happens when i change the loss calculation method to be something custom-like, just as an exampleloss_evaluated = model(xs) * + entropy_loss where entopy_loss is a value computed from the output lets say (the error happens without it too) note that the [5] comes from the batch size , i read here a solution that i should write my own torch layer instead of using the existing one, but i am not really sure if that would fix this gradient computing problem. Thanks!

Platform info: Linux-6.10.11-100.fc39.x86_64-x86_64-with-glibc2.38
Python version: 3.12.4
Numpy version: 1.26.4
Scipy version: 1.13.1
Installed devices:

    lightning.qubit (PennyLane_Lightning-0.37.0)
    qiskit.aer (PennyLane-qiskit-0.37.0)
    qiskit.basicaer (PennyLane-qiskit-0.37.0)
    qiskit.basicsim (PennyLane-qiskit-0.37.0)
    qiskit.ibmq (PennyLane-qiskit-0.37.0)
    qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.37.0)
    qiskit.ibmq.sampler (PennyLane-qiskit-0.37.0)
    qiskit.remote (PennyLane-qiskit-0.37.0)
    default.clifford (PennyLane-0.37.0)
    default.gaussian (PennyLane-0.37.0)
    default.mixed (PennyLane-0.37.0)
    default.qubit (PennyLane-0.37.0)
    default.qubit.autograd (PennyLane-0.37.0)
    default.qubit.jax (PennyLane-0.37.0)
    default.qubit.legacy (PennyLane-0.37.0)
    default.qubit.tf (PennyLane-0.37.0)
    default.qubit.torch (PennyLane-0.37.0)
    default.qutrit (PennyLane-0.37.0)
    default.qutrit.mixed (PennyLane-0.37.0)
    default.tensor (PennyLane-0.37.0)
    null.qubit (PennyLane-0.37.0)
    lightning.gpu (PennyLane_Lightning_GPU-0.37.0)
    cirq.mixedsimulator (PennyLane-Cirq-0.36.0.post0)
    cirq.pasqal (PennyLane-Cirq-0.36.0.post0)
    cirq.qsim (PennyLane-Cirq-0.36.0.post0)
    cirq.qsimh (PennyLane-Cirq-0.36.0.post0)
    cirq.simulator (PennyLane-Cirq-0.36.0.post0)
    qrack.simulator (pennylane-qrack-0.9.1)

Hi @Moustafa ,

I cannot guarantee that it will work but I strongly suggest that you follow the indication in my post from August 5th here. This should fix your first error. I would recommend starting with the code I shared and changing one element at a time instead of several at once. This can make it easier to identify and solve issues as they arise. Once you follow this approach of creating your own hybrid model let me know if you still get stuck on other issues.

Finally, I strongly suggest upgrading your PennyLane version to v0.38 (the latest one at the moment). I hope this helps!