Update README.md

README.md

@@ -70,6 +70,24 @@ print(tokenizer.decode(outputs[0]))
 ```
 
 
+#### Running the model using Flax on a single GPU / TPU 
+
+
+```python
+from transformers import AutoTokenizer, FlaxGemmaForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
+model = FlaxGemmaForCausalLM.from_pretrained("google/gemma-7b-flax")
+model.params = jax.tree_map(lambda x: jax.device_put(x, jax.devices()[0]).astype(jnp.float16), flax.params)
+
+input_text = "Write me a poem about Machine Learning."
+input_ids = jax.device_put(tokenizer(input_text, return_tensors="jax"), jax.devices()[0])
+
+outputs = model.generate(**input_ids)
+print(tokenizer.decode(outputs[0][0]))
+```
+
+
 #### Running the model on a single / multi GPU