def __init__(self, input_dataset, num_workers, index): self._input_dataset = input_dataset self._structure = input_dataset._element_structure # pylint: disable=protected-access variant_tensor = ged_ops.experimental_auto_shard_dataset( self._input_dataset._variant_tensor, # pylint: disable=protected-access num_workers=num_workers, index=index, **dataset_ops.flat_structure(self)) super(_AutoShardDataset, self).__init__(input_dataset, variant_tensor)
def __init__(self, input_dataset, num_workers, index): self._input_dataset = input_dataset self._structure = input_dataset._element_structure # pylint: disable=protected-access if compat.forward_compatible(2019, 8, 3): variant_tensor = ged_ops.auto_shard_dataset( self._input_dataset._variant_tensor, # pylint: disable=protected-access num_workers=num_workers, index=index, **self._flat_structure) else: variant_tensor = ged_ops.experimental_auto_shard_dataset( self._input_dataset._variant_tensor, # pylint: disable=protected-access num_workers=num_workers, index=index, **self._flat_structure) super(_AutoShardDataset, self).__init__(input_dataset, variant_tensor)