Do not assume 8 devices in JAX (#154)

- Do not assume 8 devices in JAX (e124bbdca2dab1af0cdce19d575f8043eab9341e)


Co-authored-by: Pedro Cuenca <pcuenq@users.noreply.huggingface.co>

README.md

@@ -154,7 +154,7 @@ prompt_ids = pipeline.prepare_inputs(prompt)
 
 # shard inputs and rng
 params = replicate(params)
-prng_seed = jax.random.split(prng_seed, 8)
+prng_seed = jax.random.split(prng_seed, num_samples)
 prompt_ids = shard(prompt_ids)
 
 images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images
@@ -187,7 +187,7 @@ prompt_ids = pipeline.prepare_inputs(prompt)
 
 # shard inputs and rng
 params = replicate(params)
-prng_seed = jax.random.split(prng_seed, 8)
+prng_seed = jax.random.split(prng_seed, num_samples)
 prompt_ids = shard(prompt_ids)
 
 images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images