diff --git a/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala b/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala
index 2d9dc21b..f6420014 100644
--- a/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala
+++ b/language_generators/scala-defref-extractor/src/main/scala/io/bazeltools/buildgen/javadefref/JavaSourceEntityExtractor.scala
@@ -61,7 +61,13 @@ object JavaSourceEntityExtractor {
 
     // TODO there needs to be some way to configure this
     val generators: Set[String] =
-      Set("InputTransformDef", "FeatureEncoderDef", "DataTransformDef")
+      Set(
+        "InputTransformDef",
+        "FeatureEncoderDef",
+        "DataTransformDef",
+        "TensorEncoderDef",
+        "InputSelectorDef"
+      )
 
     val topLevelDefsTypes: SortedSet[Entity] =
       compUnit