Update TF dataset to 2.8 (NVIDIA#3442)

* Update TF dataset to 2.8 Signed-off-by: Albert Wolant <[email protected]>
lgtm-migrator · Oct 25, 2021 · 64b79a6 · 64b79a6
1 parent 9150676
commit 64b79a6
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/dali_tf_plugin/dali_dataset_op.cc b/dali_tf_plugin/dali_dataset_op.cc
@@ -138,10 +138,20 @@ class DALIDatasetOp::Dataset : public DatasetBase {
     return "DALI::DatasetOp()::Dataset";
   }
 
+#if TF_MAJOR_VERSION == 2 && TF_MINOR_VERSION >= 8
+
+  tensorflow::int64 CardinalityInternal() const override {
+    return data::kInfiniteCardinality;
+  }
+
+#else
+
   tensorflow::int64 Cardinality() const override {
     return data::kInfiniteCardinality;
   }
 
+#endif
+
  protected:
   PipelineDef pipeline_def_;
   std::vector<PartialTensorShape> shapes_;