diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 68cef6b4a4c..e384b949d51 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -116,9 +116,9 @@ def is_transformer_engine_available():
 
 def is_transformer_engine_mxfp8_available():
     if _is_package_available("transformer_engine", "transformer-engine"):
-        import transformer_engine.pytorch as te
+        from transformer_engine.pytorch.fp8 import check_mxfp8_support
 
-        return te.fp8.check_mxfp8_support()[0]
+        return check_mxfp8_support()[0]
     return False
 
 
diff --git a/src/accelerate/utils/transformer_engine.py b/src/accelerate/utils/transformer_engine.py
index 6a950d131dd..31b61364fe2 100644
--- a/src/accelerate/utils/transformer_engine.py
+++ b/src/accelerate/utils/transformer_engine.py
@@ -154,9 +154,9 @@ def apply_fp8_autowrap(model, fp8_recipe_handler):
 
     else:
         import transformer_engine.common.recipe as te_recipe
-        import transformer_engine.pytorch as te
+        from transformer_engine.pytorch.fp8 import check_mxfp8_support
 
-        is_fp8_block_scaling_available, message = te.fp8.check_mxfp8_support()
+        is_fp8_block_scaling_available, message = check_mxfp8_support()
 
     kwargs = fp8_recipe_handler.to_kwargs() if fp8_recipe_handler is not None else {}
     if "fp8_format" in kwargs: