changes

Author: JacobSzwejbka

CMakeLists.txt

@@ -812,13 +812,11 @@ if(EXECUTORCH_BUILD_PYBIND)
     # pybind training
     pybind11_add_module(_training_lib SHARED extension/training/pybindings/_training_lib.cpp)
 
-    target_compile_definitions(
-      _training_lib PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=_training_lib
-    )
     target_include_directories(_training_lib PRIVATE ${TORCH_INCLUDE_DIRS})
     target_compile_options(_training_lib PUBLIC ${_pybind_compile_options})
     target_link_libraries(_training_lib PRIVATE ${_pybind_training_dep_libs})
 
+    message(STATUS "Installing Training Lib")
     install(TARGETS _training_lib
             LIBRARY DESTINATION executorch/extension/training/pybindings
     )

install_executorch.py

@@ -80,9 +80,10 @@ def main(args):
                     )
                 if pybind_arg == "training":
                     CMAKE_ARGS += " -DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON"
+                    os.environ["EXECUTORCH_BUILD_TRAINING"] = "ON"
                 else: 
-                    EXECUTORCH_BUILD_PYBIND = "ON"
                     CMAKE_ARGS += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"
+                EXECUTORCH_BUILD_PYBIND = "ON"
 
     if args.clean:
         clean()

setup.py

@@ -51,6 +51,7 @@
 import platform
 import re
 import sys
+import logging
 
 # Import this before distutils so that setuptools can intercept the distuils
 # imports.
@@ -88,7 +89,7 @@ def pybindings(cls) -> bool:
 
     @classmethod
     def training(cls) -> bool:
-        return cls._is_env_enabled("EXECUTORCH_BUILD_TRAINING", default=True)
+        return cls._is_env_enabled("EXECUTORCH_BUILD_TRAINING", default=False)
 
     @classmethod
     def llama_custom_ops(cls) -> bool:
@@ -580,6 +581,7 @@ def run(self):
                 "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON",
             ]
             if ShouldBuild.training():
+                logging.info("Building training foobar")
                 cmake_args += [
                     "-DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON",
                 ]
@@ -686,7 +688,6 @@ def get_ext_modules() -> List[Extension]:
             )
         )
         if ShouldBuild.training():
-
             ext_modules.append(
                 # Install the prebuilt pybindings extension wrapper for training
                 BuiltExtension(