update example (sanchit-gandhi#78)

forcs-ai · May 5, 2023 · 0d3bc54 · 0d3bc54
1 parent e2dddc9
commit 0d3bc54
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/benchmarks/run_pmap.py b/benchmarks/run_pmap.py
@@ -3,6 +3,7 @@
 import jax
 import jax.numpy as jnp
 from datasets import concatenate_datasets, load_dataset
+from flax import jax_utils
 from flax.training.common_utils import shard
 from transformers import FlaxWhisperForConditionalGeneration, WhisperProcessor
 
@@ -17,7 +18,7 @@
     dtype=jnp.bfloat16,
 )
 
-params = model.to_bf16(params)
+params = jax_utils.replicate(params)
 
 
 def generate_fn(batch):