diff --git a/src/fairchem/core/common/relaxation/__init__.py b/src/fairchem/core/common/relaxation/__init__.py
index 9fea399f0..068eaf718 100644
--- a/src/fairchem/core/common/relaxation/__init__.py
+++ b/src/fairchem/core/common/relaxation/__init__.py
@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+from .ml_relaxation import ml_relax
 from .optimizers.optimizable import OptimizableBatch, OptimizableUnitCellBatch
 
-__all__ = ["OptimizableBatch", "OptimizableUnitCellBatch"]
+__all__ = ["ml_relax", "OptimizableBatch", "OptimizableUnitCellBatch"]