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).
See the original Flux issue and part of the solution here.