Documentation fixes (#248)

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Documentation fixes (#248)
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
215dfe7e · Kirthi Shankar Sivamani · GitHub · 156a074a · 215dfe7e · 215dfe7e
Unverified Commit 215dfe7e authored May 26, 2023 by Kirthi Shankar Sivamani Committed by GitHub May 26, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

docs/api/jax.rst docs/api/jax.rst +1 -1

transformer_engine/pytorch/export.py transformer_engine/pytorch/export.py +1 -0

No files found.
--- a/docs/api/jax.rst
+++ b/docs/api/jax.rst
@@ -8,7 +8,7 @@ Jax

 .. autoapiclass:: transformer_engine.jax.MajorShardingType
 .. autoapiclass:: transformer_engine.jax.ShardingType
-.. autoapiclass:: transformer_engine.jax.TransformerLayerType
+.. autoapiclass:: transformer_engine.jax.flax.TransformerLayerType
 .. autoapiclass:: transformer_engine.jax.ShardingResource(dp_resource=None, tp_resource=None)



--- a/transformer_engine/pytorch/export.py
+++ b/transformer_engine/pytorch/export.py
@@ -19,6 +19,7 @@ def onnx_export(
        with onnx_export(enabled=True):
            torch.onnx.export(model)

+    Parameters
    ----------
    enabled: bool, default = `False`
             whether or not to enable export