Hi @CatalinaAlbornoz,

Thanks for looking into my query. Below is my code example:

# imports

```
import pennylane as qml
from pennylane import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn import datasets
```

# dataset loading

`x_train, x_test = datasets.load_digits().data, datasets.load_digits().target`

# splitting train and test data

```
x_train, x_test, y_train, y_test = train_test_split(
x_train, x_test, test_size=0.25, random_state=42)
```

# creating qnode

```
n_qubits = 6
dev = qml.device("default.qubit", wires=n_qubits)
@qml.qnode(dev)
def qnode(inputs, w1):
qml.templates.AmplitudeEmbedding(inputs, wires=range(n_qubits), pad_with=2, normalize = True)
for i in range(n_qubits):
qml.RY(w1[i], wires = i)
return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]
```

# defining weight shapes for quantum gates

`weight_shapes = {"w1": (1,6)}`

# qnode to keras layer

`qlayer = qml.qnn.KerasLayer(qnode, weight_shapes, output_dim=n_qubits)`

# creating a hybrid NN

```
clayer = tf.keras.layers.Dense(10, activation="softmax")
model = tf.keras.models.Sequential([qlayer, clayer])
```

# compiling the NN model

```
opt = tf.keras.optimizers.SGD(learning_rate=0.2)
model.compile(opt, loss="sparse_categorical_crossentropy", metrics=["accuracy"])
```

# start training the model

`fitting = model.fit(x_train, y_train, epochs=1, batch_size=16, validation_data=(x_test, y_test))`

The above results in the following error:

`InvalidArgumentError: slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: sequential_20/keras_layer_21/strided_slice/`

Now the issue I am facing here is regarding the `weight_shapes`

. since I am using single parametrized gate `RY`

means it should have a single dimension weight (as per my understanding) and hence the first argument in weight_shapes dictionary is `1`

and since I am using 6 qubits the second argument is `6`

. which means `6`

weights of dimension 1 (correct me if I am wrong)? but it gives the error mentioned above.

If I set the `weight_shapes = {"w1": (1)}`

and not use `w1[i]`

in the qnode where I pass the weights to `RY`

the code works fine but in that case, I think we are only training the `6 RY`

gates with same weights which is, in a way, same like training a single weight or single RY gate, whereas all 6 RY gates should have different weights, right?

I hope I made my question clear. Any help would be appreciated.

Thanks