Call dataset cardinality op on the dataset device.

Dataset cardinality op should be placed on the same host device as the dataset, otherwise TF will try to encode/decode the dataset to another host device, which is not possible.

PiperOrigin-RevId: 339942622
Change-Id: Ifbeb0aabdfc35f6c0e1248d69a8d22727e19c4a2

tensorflow/python/distribute/input_lib.py

@@ -2146,7 +2146,8 @@ def _enable_get_next_as_optional(strategy, dataset):
     # support. Currently the logic only applies to the case that distributed
     # dataset is created in eager mode, as we need to evaluate the dataset
     # cardinality.
-    return cardinality.cardinality(dataset).numpy() != cardinality.INFINITE
+    with ops.device(dataset._variant_tensor.device):  # pylint: disable=protected-access
+      return dataset.cardinality().numpy() != cardinality.INFINITE
 
   return not _is_statically_shaped(
       dataset.element_spec) or strategy.extended._in_multi_worker_mode()  # pylint: disable=protected-access