Update TF dataset API usage to align with 2.13rc (#4707)

- related to tensorflow/tensorflow@1d795c6 Signed-off-by: Janusz Lisiecki <[email protected]>
NVIDIA · Mar 14, 2023 · 163462f · 163462f
1 parent 0882348
commit 163462f
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/dali_tf_plugin/dali_dataset_op.cc b/dali_tf_plugin/dali_dataset_op.cc
@@ -142,7 +142,13 @@ class DALIDatasetOp::Dataset : public DatasetBase {
     return "DALI::DatasetOp()::Dataset";
   }
 
-#if TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 8
+#if TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 13
+
+  int64_t CardinalityInternal(CardinalityOptions options) const override {
+    return data::kInfiniteCardinality;
+  }
+
+#elif TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 8
 
   tensorflow::int64 CardinalityInternal() const override {
     return data::kInfiniteCardinality;
@@ -156,6 +162,7 @@ class DALIDatasetOp::Dataset : public DatasetBase {
 
 #endif
 
+
  protected:
   PipelineDef pipeline_def_;
   std::vector<PartialTensorShape> shapes_;
@@ -310,7 +317,11 @@ class DALIDatasetOp::Dataset::Iterator : public DatasetIterator<Dataset> {
             " for output ", i);
 
         if (dataset()->fail_on_device_mismatch_) {
+#if TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 13
+          return Status(absl::StatusCode::kInternal, msg);
+#else
           return Status(tensorflow::error::Code::INTERNAL, msg);
+#endif
         }
       }
     }