How to handle Batch dimension of classical layers in Hybrid quantum ML model through torchlayer

PennyLane Help
1

how do I go with the given example in the section ‘Creating non-sequential models’ form the demo Turning quantum nodes into Torch Layers | PennyLane Demos when batch dimension is also present as input from the classical layer.

class HybridModel(torch.nn.Module):
def init(self):
super().init()
self.clayer_1 = torch.nn.Linear(2, 4)
self.qlayer_1 = qml.qnn.TorchLayer(qnode, weight_shapes)
self.qlayer_2 = qml.qnn.TorchLayer(qnode, weight_shapes)
self.clayer_2 = torch.nn.Linear(4, 2)
self.softmax = torch.nn.Softmax(dim=1)

def forward(self, x):
    x = self.clayer_1(x)
    x_1, x_2 = torch.split(x, 2, dim=1)
    x_1 = self.qlayer_1(x_1)
    x_2 = self.qlayer_2(x_2)
    x = torch.cat([x_1, x_2], axis=1)
    x = self.clayer_2(x)
    return self.softmax(x)

model = HybridModel()

Also, consider weight-shapes = {“weights”: (wire*2,)}

# Put code here


If you want help with diagnosing an error, please put the full error message below:

Put full error message here

And, finally, make sure to include the versions of your packages. Specifically, show us the output of `qml.about()`.