Whenever I use the following code in “forward” method of class B, it does not work for multi-instances:
for i in range(10):
    out = self.a[i](X[i])

My mean is that suppose the input array be in shape (batch_size, 5, 10) and batch_size = 3
then, when I input the input array into instance of class B, it will rise an error of array index out of bound.
b = B()
y_hat = b(input_array[index : index + batch_size])

Would you please help me to know how to work with batches in this case?