Skip to content

Commit 320b81d

Browse files
Update README to list JAX NGC container (NVIDIA#1756)
Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
1 parent 26ecb2f commit 320b81d

1 file changed

Lines changed: 8 additions & 0 deletions

File tree

README.rst

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,12 +170,20 @@ Docker (Recommended)
170170
^^^^^^^^^^^^^^^^^^^
171171
The quickest way to get started with Transformer Engine is by using Docker images on
172172
`NVIDIA GPU Cloud (NGC) Catalog <https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch>`_.
173+
174+
173175
For example to use the NGC PyTorch container interactively,
174176

175177
.. code-block:: bash
176178
177179
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.04-py3
178180
181+
For example to use the NGC JAX container interactively,
182+
183+
.. code-block:: bash
184+
185+
docker run --gpus all -it --rm nvcr.io/nvidia/jax:25.04-py3
186+
179187
Where 25.04 (corresponding to April 2025 release) is the container version.
180188

181189
**Benefits of using NGC containers:**

0 commit comments

Comments
 (0)