[BUG] Shot-Batching does not work with TorchLayers

[aaronmkts]

Expected behavior

Shot batching is where you can run a Qnode for x amount of shots, 'batch_size' number of times. A simple example is here:

import pennylane as qml
import torch
from math import pi
from torch.nn.functional import binary_cross_entropy

batch_size = 5

n_qubits = 10
depth = 4
wires = range(n_qubits)

dev = qml.device("default.qubit")

@qml.qnode(dev, diff_method="parameter-shift")
def circuit(q_weights_y, q_weights_z):
    for i in range(depth):
        for y in range(n_qubits):
            qml.RY(q_weights_y[i][y], wires=y)
            qml.RZ(q_weights_z[i][y], wires=y)
        for y in range(n_qubits - 1):
            qml.CNOT(wires=[y, y + 1])
    return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]

q_weights_y = torch.rand((depth, n_qubits), requires_grad=True, dtype=torch.float64) * 2 * pi - pi
q_weights_z = torch.rand((depth, n_qubits), requires_grad=True, dtype=torch.float64) * 2 * pi - pi
q_weights_y.retain_grad()
q_weights_z.retain_grad()

# # Probability distribution sampled from PennyLane circuit
samples = circuit(q_weights_y, q_weights_z, shots=[1] * batch_size)

Where 1 shot of the circuit is performed 5 times.

Actual behavior

Now when trying to define a similar circuit that works with the TorchLayer interface the above is no longer possible.

from pennylane.qnn import TorchLayer
import pennylane as qml

    batch_size = 2
    n_qubits = 3
    depth = 4

    dev = qml.device("default.qubit", shots= [1] *batch_size)

    def circuit():

        @qml.qnode(dev, diff_method="parameter-shift", interface = 'torch')
        def _qnode(inputs, q_weights_y, q_weights_z):
            for i in range(depth):
                for y in range(n_qubits):
                    qml.RY(q_weights_y[i][y], wires=y)
                    qml.RZ(q_weights_z[i][y], wires=y)
                for y in range(n_qubits - 1):
                    qml.CNOT(wires=[y, y + 1])
            return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]
        
        q_weight_shapes = {"q_weights_y": (depth, n_qubits),
                           "q_weights_z": (depth, n_qubits)}
        
        q_layer = TorchLayer(_qnode, q_weight_shapes)

        return q_layer

    
    q_layer = circuit()
    inputs = torch.tensor([]) #dummy input
    output = q_layer(inputs)
    print(output)
    ```
    The above code will always throw up an error.

### Additional information

I believe the error originates from within the `TorchLayer` class, specifically the `_evaluate_qnode` function

def _evaluate_qnode(self, x):
"""Evaluates the QNode for a single input datapoint.

    Args:
        x (tensor): the datapoint

    Returns:
        tensor: output datapoint
    """
    kwargs = {
        **{self.input_arg: x},
        **{arg: weight.to(x) for arg, weight in self.qnode_weights.items()},
    }
    res = self.qnode(**kwargs)

    if isinstance(res, torch.Tensor):
        return res.type(x.dtype)

    if len(x.shape) > 1:
        res = [torch.reshape(r, (x.shape[0], -1)) for r in res]

    return torch.hstack(res).type(x.dtype)
    ```
    
   From what I can see the variable `res` does in fact contain the correct information for the computation. Just before the return line in this function `res` is in fact the same output that I would expect from the working example provided. However the `torch.hstack(res).type(x.dtype)` line will not work as indeed there is a list of tensors.

Source code

No response

Tracebacks

in _evaluate_qnode
    return torch.hstack(res).type(x.dtype)
           ^^^^^^^^^^^^^^^^^
TypeError: expected Tensor as element 0 in argument 0, but got list

System information

Name: PennyLane
Version: 0.34.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane

Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           macOS-14.3.1-arm64-arm-64bit
Python version:          3.11.4
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.34.0)
- default.gaussian (PennyLane-0.34.0)
- default.mixed (PennyLane-0.34.0)
- default.qubit (PennyLane-0.34.0)
- default.qubit.autograd (PennyLane-0.34.0)
- default.qubit.jax (PennyLane-0.34.0)
- default.qubit.legacy (PennyLane-0.34.0)
- default.qubit.tf (PennyLane-0.34.0)
- default.qubit.torch (PennyLane-0.34.0)
- default.qutrit (PennyLane-0.34.0)
- null.qubit (PennyLane-0.34.0)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.

[albi3ro]

Thanks for opening this issue @aaronmkts . This does indeed seem to be a bug. We will look into it and come up with a fix when we can.