commit change

mosaicml · Jan 26, 2024 · 1a0685f · 1a0685f
1 parent 322e90f
commit 1a0685f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/composer/core/precision.py b/composer/core/precision.py
@@ -78,7 +78,8 @@ def get_precision_context(precision: Union[str, Precision],
                 }
             fp8_recipe = DelayedScaling(**precision_config)
             with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
-                yield
+                with te.onnx_export(enabled=True):
+                    yield
         else:
             if te_installed:
                 raise RuntimeError('AMP_FP8 precision is used but current device does not support it.')