diff --git a/.ci/scripts/unittest-buck2.sh b/.ci/scripts/unittest-buck2.sh
index 340f7438f02..295295340fc 100755
--- a/.ci/scripts/unittest-buck2.sh
+++ b/.ci/scripts/unittest-buck2.sh
@@ -15,7 +15,8 @@ buck2 query "//backends/apple/... + //backends/arm: + //backends/arm/debug/... +
 //backends/arm/_passes/... + //backends/arm/runtime/... + //backends/arm/tosa/... \
 + //backends/example/... + \
 //backends/mediatek/... + //backends/transforms/... + \
-//backends/xnnpack/... + //configurations/... + //extension/flat_tensor: + \
+//backends/xnnpack/... + //codegen/tools/... + \
+//configurations/... + //extension/flat_tensor: + \
 //extension/llm/runner: + //kernels/aten/... + //kernels/optimized/... + \
 //kernels/portable/... + //kernels/quantized/... + //kernels/test/... + \
 //runtime/... + //schema/... + //test/... + //util/..."
@@ -38,3 +39,6 @@ for op in "build" "test"; do
           $BUILDABLE_KERNELS_PRIM_OPS_TARGETS //runtime/backend/... //runtime/core/... \
           //runtime/executor: //runtime/kernel/... //runtime/platform/...
 done
+
+# Build only without testing
+buck2 build //codegen/tools/... # Needs torch for testing which we don't have in our OSS buck setup.
diff --git a/codegen/tools/targets.bzl b/codegen/tools/targets.bzl
index d594b7178b8..c11982409f0 100644
--- a/codegen/tools/targets.bzl
+++ b/codegen/tools/targets.bzl
@@ -17,10 +17,8 @@ def define_common_targets(is_fbcode = False):
         ],
         deps = [
             "//executorch/codegen:gen_lib",
-        ] + ([] if runtime.is_oss else select({
-            "DEFAULT": [],
-            "ovr_config//os:linux": ["//executorch/codegen/tools:selective_build"],  # TODO(larryliu0820) :selective_build doesn't build in OSS yet
-        })),
+            "//executorch/codegen/tools:selective_build",
+        ],
     )
 
     runtime.python_binary(
@@ -29,7 +27,7 @@ def define_common_targets(is_fbcode = False):
         deps = [
             ":gen_oplist_lib",
         ],
-        preload_deps = [] if runtime.is_oss else ["//executorch/codegen/tools:selective_build"],  # TODO(larryliu0820) :selective_build doesn't build in OSS yet
+        preload_deps = ["//executorch/codegen/tools:selective_build"],
         package_style = "inplace",
         visibility = [
             "//executorch/...",
@@ -196,27 +194,27 @@ def define_common_targets(is_fbcode = False):
         _is_external_target = True,
     )
 
-    if not runtime.is_oss:
-        runtime.cxx_python_extension(
-            name = "selective_build",
-            srcs = [
-                "selective_build.cpp",
-            ],
-            base_module = "executorch.codegen.tools",
-            types = ["selective_build.pyi"],
-            preprocessor_flags = [
-                "-DEXECUTORCH_PYTHON_MODULE_NAME=selective_build",
-            ],
-            deps = [
-                "//executorch/runtime/core:core",
-                "//executorch/schema:program",
-            ],
-            external_deps = [
-                "pybind11",
-            ],
-            use_static_deps = True,
-            visibility = ["//executorch/codegen/..."],
-        )
+    
+    runtime.cxx_python_extension(
+        name = "selective_build",
+        srcs = [
+            "selective_build.cpp",
+        ],
+        base_module = "executorch.codegen.tools",
+        types = ["selective_build.pyi"],
+        preprocessor_flags = [
+            "-DEXECUTORCH_PYTHON_MODULE_NAME=selective_build",
+        ],
+        deps = [
+            "//executorch/runtime/core:core",
+            "//executorch/schema:program",
+        ],
+        external_deps = [
+            "pybind11",
+        ],
+        use_static_deps = True,
+        visibility = ["//executorch/codegen/..."],
+    )
 
 
     # TODO(larryliu0820): This is a hack to only run these two on fbcode. These targets depends on exir which is only available in fbcode.
@@ -255,10 +253,12 @@ def define_common_targets(is_fbcode = False):
             ],
         )
 
+    if runtime.is_oss or is_fbcode:
+        # Doesn't work on xplat. But works on fbcode and OSS.
         runtime.python_test(
-            name = "test_selective_build",
+            name = "test_tools_selective_build",
             srcs = [
-                "test/test_selective_build.py",
+                "test/test_tools_selective_build.py",
             ],
             package_style = "inplace",
             visibility = [
diff --git a/codegen/tools/test/test_selective_build.py b/codegen/tools/test/test_tools_selective_build.py
similarity index 100%
rename from codegen/tools/test/test_selective_build.py
rename to codegen/tools/test/test_tools_selective_build.py
diff --git a/extension/pytree/TARGETS b/extension/pytree/TARGETS
index 5f09772ee38..79c4d2b8b3f 100644
--- a/extension/pytree/TARGETS
+++ b/extension/pytree/TARGETS
@@ -15,9 +15,11 @@ runtime.cxx_python_extension(
     ],
     base_module = "executorch.extension.pytree",
     deps = [
-        "fbsource//third-party/pybind11:pybind11",
         ":pytree",
     ],
+    external_deps = [
+        "pybind11",
+    ],
 )
 
 runtime.cxx_python_extension(
@@ -27,9 +29,11 @@ runtime.cxx_python_extension(
     ],
     base_module = "executorch.extension.pytree",
     deps = [
-        "fbsource//third-party/pybind11:pybind11",
         ":pytree",
     ],
+    external_deps = [
+        "pybind11",
+    ],
 )
 
 runtime.python_library(
diff --git a/pytest.ini b/pytest.ini
index 100c47aed50..7c722d50e29 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -31,6 +31,7 @@ addopts =
 
     # codegen
     codegen/test
+    codegen/tools/test/test_tools_selective_build.py
 
     # devtools
     devtools/
diff --git a/shim_et/xplat/executorch/build/runtime_wrapper.bzl b/shim_et/xplat/executorch/build/runtime_wrapper.bzl
index 46101ab36db..d03a3f71b3b 100644
--- a/shim_et/xplat/executorch/build/runtime_wrapper.bzl
+++ b/shim_et/xplat/executorch/build/runtime_wrapper.bzl
@@ -346,6 +346,7 @@ def _python_binary(*args, **kwargs):
 
 def _python_test(*args, **kwargs):
     _patch_kwargs_common(kwargs)
+    _remove_caffe2_deps(kwargs)
     env.python_test(*args, **kwargs)
 
 def get_oss_build_kwargs():