Wrapping a Flux FCN model with MLJFlux

General Usage
,
1

Cheers,

I wonder if it is possible to wrap a fully convolutional network written in Flux (e.g. a U-Net) with MLJFlux, and benefit from the high-level interface from MLJ. FCN outputs are typically multi-dimensional array (as in the U-Net), rather than a vector (one-dimension array, as in a classifier) as stated in this page:

The object returned by chain(x) must be an AbstractFloat vector of length n_out

If MLJ is indeed not feasible for this task, please kindly advise on alternatives that could provide high-level interface (e.g. cross-validation) for such Flux models.

Thanks in advance.