cannot find good documentation of this anywhere
as far as I can tell, it must be an array of arrays where each array is of only one dimension
It can be anything that’s iterable, where each iterated element is also iterable (so an array of arrays works), and the iterated elements will just be splatted into the loss function. If you look at the definition of train!, this should be pretty obvious: https://github.com/FluxML/Flux.jl/blob/b47238eb7447311e009a2413a317f7216fe54d12/src/optimise/train.jl#L69
I took a look at that but it was not obvious to me, likely since I am still pretty new to Julia