Key-value pair in data does not have same shape: (2,), (500, 2)

General
1

inference = ed.SGLD({pi: qpi, mu: qmu, sigmasq: qsigmasq, z:qz},
data={x: x_train})