diff --git a/setup.py b/setup.py
index afb67496136f..6976e79b539d 100644
--- a/setup.py
+++ b/setup.py
@@ -116,6 +116,7 @@ def cmake_build(self, cmake_build_dir):
             f"-DPython_EXECUTABLE={sys.executable}",
             f"-DPython_FIND_VIRTUALENV=ONLY",
             f"-DMLIR_ENABLE_BINDINGS_PYTHON=ON",
+            "-DMLIR_BINDINGS_PYTHON_NB_DOMAIN=torch_mlir",
             f"-DLLVM_TARGETS_TO_BUILD=host",
             f"-DLLVM_ENABLE_ZSTD=OFF",
             # Optimization options for building wheels.