3D Conv Training Error

Hi there,

I encountered an error while setting up a 3D CNN that I dont understand. The problems boils down to this:

While in 2D

using Flux

a, b = rand(5,10,2,10), rand(5,10,2,10)
m= Chain(Conv((3,3),2=>4,pad=(1,1)), Conv((3,3), 4=>2, pad=(1,1)))
loss(x,y) = Flux.mse(m(x),y)
@time Flux.train!(loss, params(m), [[a,b]], ADAM(0.01))

works without problems, in 3D

a2, b2 = rand(5,10,10,2,10), rand(5,10,10,2,10)
m2= Chain(Conv((3,3,3),2=>4,pad=(1,1,1)), Conv((3,3,3), 4=>2, pad=(1,1,1)))
loss2(x,y) = Flux.mse(m2(x),y)
@time Flux.train!(loss2, params(m2), [[a2,b2]], ADAM(0.01))

will result in a dimension mismatch on the training (back/gradient operator) even though loss2(a2,b2) works.

Am I doing something wrong here or is this a bug?

Best,
Max