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

I’m trying to run the file found here: model-zoo/vae_mnist.jl at master · FluxML/model-zoo · GitHub on newer versions of Flux, CUDA.

CUDA = v3.4.2
Flux = v0.12.6

Do you know what could be causing errors with getting the loss/gradient? It seems odd that this step is broken.

When I run the example with the latest versions of Flux, CUDA on a gpu, I get the error “ERROR: this intrinsic must be compiled to be called” when the model loss is being calculated.

When I run the example on the cpu, the loss diverges.

Updating Zygote to the latest version on github fixes the GPU error, but the loss still diverges.

Zygote v0.6.22 fixed this issue