Minibatch Theano: shape mismatch:

I am using toy data used in this posthttps://twiecki.io/blog/2018/08/13/hierarchical_bayesian_neural_network/.

I used the following code for variation inference:

Xss = theano.shared(Xs_train)
Yss = theano.shared(Ys_train)
hbnn = construct_hierarchical_nn(Xss, Yss)

minibatch_x = pm.Minibatch(Xs_train, batch_size=[1,50,2])
minibatch_y = pm.Minibatch(Ys_train, batch_size=[1,50])

with hbnn:
advi_fit = pm.fit(method=pm.ADVI(), n=10000,
more_replacements={Xss: minibatch_x, Yss:minibatch_y})

I got the following error:

   ValueError: Shape mismatch: batch sizes unequal. x.shape is (1, 50, 2), y.shape is (4, 2, 5).
  **Apply node that caused the error: BatchedDot(Minibatch, Elemwise{Composite{((i0 * i1) + i2)}}.0)**

** Toposort index: 126**
** Inputs types: [TensorType(float64, 3D), TensorType(float64, 3D)]**
** Inputs shapes: [(1, 50, 2), (4, 2, 5)]**
** Inputs strides: [(800, 16, 8), (80, 40, 8)]**
** Inputs values: [‘not shown’, ‘not shown’]**
** Outputs clients: [[Elemwise{Tanh}(0, 0)]]**

Any help on how I can fix this!!