[Flax] Solve problem with VAE (#574)

2a8477de · Patrick von Platen · GitHub · bf5ca036 · 2a8477de
Unverified Commit 2a8477de authored Sep 19, 2022 by Patrick von Platen Committed by GitHub Sep 19, 2022
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

src/diffusers/models/vae_flax.py src/diffusers/models/vae_flax.py +2 -1

No files found.
--- a/src/diffusers/models/vae_flax.py
+++ b/src/diffusers/models/vae_flax.py
@@ -600,7 +600,8 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin):
            hidden_states = posterior.latent_dist.sample(rng)
        else:
            hidden_states = posterior.latent_dist.mode()
-        hidden_states = self.decode(hidden_states, return_dict=return_dict).sample
+
+        sample = self.decode(hidden_states, return_dict=return_dict).sample

        if not return_dict:
            return (sample,)