examples/jax/mnist/test_single_gpu_mnist.py · 5b6ef0545ed14497ff296d18c3618ace4ba391db · OpenDAS / TransformerEngine · GitLab

Find file Blame History Permalink

[JAX] Flax with compute dtype inferred from input dtype. (#1485) · 6673f165
Phuong Nguyen authored Feb 18, 2025
```
flax module with compute dtype inferred from the inputs
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
```
6673f165

test_single_gpu_mnist.py 10.2 KB