examples/jax/mnist/test_single_gpu_mnist.py · 838345eba4fdd2a169dd9e087d39c30a360e684a · OpenDAS / TransformerEngine

[Common] Remove CheckTensor if the workspace is empty in cast_transpose_fused (#931) · 4a4f05da

Phuong Nguyen authored Jun 15, 2024



* rm tensor check if the workspace is empty
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add trust_remote=true for load_dataset() in the mnist test
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

4a4f05da

test_single_gpu_mnist.py 10 KB

Replace test_single_gpu_mnist.py