Fix for FluxML#301 according to discourse

This fixes FluxML#301 by following advice from https://discourse.julialang.org/t/mutating-arrays-is-not-supported-error-when-running-fluxml-vae-mnist-jl/59890
jstrube · May 8, 2021 · 0436b34 · 0436b34
1 parent b71bf7b
commit 0436b34
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/vision/vae_mnist/vae_mnist.jl b/vision/vae_mnist/vae_mnist.jl
@@ -55,6 +55,7 @@ function reconstuct(encoder, decoder, x, device)
     μ, logσ, decoder(z)
 end
 
+Flux.Zygote.@nograd Flux.params
 function model_loss(encoder, decoder, λ, x, device)
     μ, logσ, decoder_z = reconstuct(encoder, decoder, x, device)
     len = size(x)[end]