Skip to content

GitLab

Unverified Commit 60438389 authored by Pedro Cuenca's avatar Pedro Cuenca Committed by GitHub
Browse files

Fix OOM when using PyTorch with JAX installed. (#1795) 

Don't initialize Jax on startup.
parent 4125756e
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
src/diffusers/schedulers/scheduling_ddpm_flax.py
View file @ 60438389
...@@ -208,7 +208,7 @@ class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin): ...@@ -208,7 +208,7 @@ class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin):
model_output: jnp.ndarray, model_output: jnp.ndarray,
timestep: int, timestep: int,
sample: jnp.ndarray, sample: jnp.ndarray,
key: jax.random.KeyArray = jax.random.PRNGKey(0), key: jax.random.KeyArray,
return_dict: bool = True, return_dict: bool = True,
) -> Union[FlaxDDPMSchedulerOutput, Tuple]: ) -> Union[FlaxDDPMSchedulerOutput, Tuple]:
""" """
......
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