See the original [Flux issue](https://github.com/FluxML/Flux.jl/issues/1596) and part of the solution [here](https://github.com/FluxML/Flux.jl/pull/1595#issuecomment-840270106).