[PyTorch] Fix ONNX/dynamo error (#562)

* fix onnx/dynamo error
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* move changes to pytorch/__init__ using try/except
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>

* Update transformer_engine/pytorch/__init__.py
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>

---------
Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
Signed-off-by: cyanguwa <8636796+cyanguwa@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>

transformer_engine/pytorch/__init__.py

@@ -29,3 +29,8 @@ from .te_onnx_extensions import (
     onnx_rmsnorm_fwd,
     onnx_rmsnorm_fwd_fp8
 )
+try:
+    import torch
+    torch._dynamo.config.error_on_nested_jit_trace = False
+except:
+    pass