Fix docstring for DistributedEmbedding. (#118)

The JAX usage section was broken.

Author: hertschuh

keras_rs/src/layers/embedding/base_distributed_embedding.py

@@ -297,6 +297,7 @@ def step(data):
             result = strategy.run(step, args=(next(iterator),))
 
     run_loop(iter(dataset))
+    ```
 
     ### Usage with JAX on TPU with SpareCore