Skip to content

GitLab

"src/vscode:/vscode.git/clone" did not exist on "c4892f1855097a68703ca2e949aca15829526958"
Unverified Commit 2a8477de authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

[Flax] Solve problem with VAE (#574)

parent bf5ca036
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
src/diffusers/models/vae_flax.py
View file @ 2a8477de
...@@ -600,7 +600,8 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin): ...@@ -600,7 +600,8 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin):
hidden_states = posterior.latent_dist.sample(rng) hidden_states = posterior.latent_dist.sample(rng)
else: else:
hidden_states = posterior.latent_dist.mode() 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: if not return_dict:
return (sample,) return (sample,)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment