Cleanup: Fix includes.

PiperOrigin-RevId: 729611238
jessejay-ch · Feb 21, 2025 · 44dc41d · 44dc41d
1 parent 62a501e
commit 44dc41d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
@@ -739,6 +739,7 @@ tf_kernel_library(
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/data:name_utils",
+        "@com_google_absl//absl/strings",
     ],
 )
 

diff --git a/tensorflow/core/kernels/data/experimental/rebatch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/rebatch_dataset_op.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <optional>
 
+#include "absl/strings/str_join.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/tensor.h"
@@ -330,7 +331,7 @@ class RebatchDatasetV2Op : public UnaryDatasetOpKernel {
           output_types_(output_types),
           output_shapes_(output_shapes),
           traceme_metadata_(
-              {{"batch_sizes", absl::StrJoin(batch_sizes, ",")}}) {
+              {{"batch_sizes", absl::StrJoin(batch_sizes_, ",")}}) {
       input_->Ref();
     }