Flux with Matrix input dimensions

Hi

Just started using flux, cant get on top of input/layer dimensions where input is a matrix

Adapting https://github.com/FluxML/model-zoo/blob/master/vision/mnist/mlp.jl to my input space but its not working

My Input set is Vector{Array{Array{Float32,2},1}}: Vector of 5x10 Matrices of Float32
Output set is an Array of 1’s and 0’s

I use flatten : xtrain = Flux.flatten(xtrain)
like the examples does to reduce to Array{Array{Float32,2},2}
I use the Dataoader like the example does:

train_data = DataLoader(xtrain, ytrain, batchsize=32, shuffle=true)
test_data = DataLoader(xtest, ytest, batchsize=32)

I use the model:
function build_model(; imgsize=(5,10,1), nclasses=2)
return Chain(
Dense(prod(imgsize), 32, relu),
Dense(32, nclasses,softmax))
end

I call @epochs args.epochs Flux.train!(loss, Flux.params(m), train_data, opt, cb = evalcb)
then the loss function gives an error
loss(x,y) = logitcrossentropy(m(x), y)

When I debug its the evaluation of m(x)
“Error: DimensionMismatch(“matrix A has dimensions (32,50), matrix B has dimensions (1,32)”)”
32 is the size of the hidden layer
50 is the size of the input matrix
32 is the size of a batch

Any help/insight will be appreciated

Thanks

Managed to solve myself by changing the input definition to 5x10xN