Error when trying to run /FluxML/model-zoo/vision/vae_mnist/vae_mnist.jl on newer versions of CUDA, Flux

Zygote v0.6.22 fixed this issue