diff --git a/CMakeLists.txt b/CMakeLists.txt
index 38d9f26..8850b8e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -39,6 +39,8 @@
 include_directories(${MLIR_INCLUDE_DIRS})
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/include)
+# Needed to build TorchMLIRExtensions.
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/externals/torch-mlir/include)
 
 function(mpact_target_includes target)
   set(_dirs
diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
index 7815708..f8d71ac 100644
--- a/python/CMakeLists.txt
+++ b/python/CMakeLists.txt
@@ -21,7 +21,14 @@
   ADD_TO_PARENT MPACTPythonSources
   SOURCES
     mpactbackend.py
-  )
+)
+
+declare_mlir_python_sources(MPACTPythonSources.SampleModels
+  ROOT_DIR "${MPACT_PYTHON_ROOT_DIR}"
+  ADD_TO_PARENT MPACTPythonSources
+  SOURCES_GLOB
+    models/*.py
+)
 
 #-------------------------------------------------------------------------------
 # Extensions
@@ -38,22 +45,10 @@
     LLVMSupport
 )
 
-declare_mlir_python_sources(MPACTPythonSources.SampleModels
-  ROOT_DIR "${MPACT_PYTHON_ROOT_DIR}"
-  ADD_TO_PARENT MPACTPythonSources
-  SOURCES_GLOB
-    models/*.py
-)
-
 #-------------------------------------------------------------------------------
 # Python Modules
 #-------------------------------------------------------------------------------
 
-# To compiler TorchMLIRPythonExtension (which registers torch dialect and
-# related passes).
-# TODO: is there a cleaner way?
-include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../externals/torch-mlir/include)
-
 set(_source_components
   MLIRPythonSources
   MLIRPythonExtension.Core
@@ -61,6 +56,7 @@
 
   # We need the FxImporter from torch-mlir
   TorchMLIRPythonSources.Importers
+  TorchMLIRPythonSources.Dialects
   TorchMLIRPythonExtensions
 
   MPACTPythonSources