message: [tf.data] Disabling eager garbage collection of function graphs as in some cases, the function graph may be collected prematurely.
diff: @@ -51,8 +51,7 @@ class _PerDeviceGenerator(dataset_ops.DatasetV2):
  def _init_func():
  return multi_device_iterator_string_handle
  - init_func_concrete = (
  - _init_func._get_concrete_function_internal_garbage_collected()) # pylint: disable=protected-access
  + init_func_concrete = _init_func._get_concrete_function_internal() # pylint: disable=protected-access
  # TODO(b/124254153): Enable autograph once the overhead is low enough.
  @function.defun(autograph=False) # Pure graph code.
  @@ -63,8 +62,7 @@ class _PerDeviceGenerator(dataset_ops.DatasetV2):
  Tout=[dtypes.string],
  f=init_func_concrete)
  - self._init_func = (
  - _remote_init_func._get_concrete_function_internal_garbage_collected()) # pylint: disable=protected-access
  + self._init_func = _remote_init_func._get_concrete_function_internal() # pylint: disable=protected-access
  self._init_captured_args = self._init_func.captured_inputs
  # TODO(b/124254153): Enable autograph once the overhead is low enough.
  @@ -86,8 +84,7 @@ class _PerDeviceGenerator(dataset_ops.DatasetV2):
  output_types=structure.get_flat_tensor_types(self._element_spec),
  output_shapes=structure.get_flat_tensor_shapes(self._element_spec))
  - next_func_concrete = (
  - _next_func._get_concrete_function_internal_garbage_collected()) # pylint: disable=protected-access
  + next_func_concrete = _next_func._get_concrete_function_internal() # pylint: disable=protected-access
  # TODO(b/124254153): Enable autograph once the overhead is low enough.
  @function.defun_with_attributes(
  @@ -101,8 +98,7 @@ class _PerDeviceGenerator(dataset_ops.DatasetV2):
  Tout=structure.get_flat_tensor_types(self._element_spec),
  f=next_func_concrete)
  - self._next_func = (
  - _remote_next_func._get_concrete_function_internal_garbage_collected()) # pylint: disable=protected-access
  + self._next_func = _remote_next_func._get_concrete_function_internal() # pylint: disable=protected-access
  self._next_captured_args = self._next_func.captured_inputs
  self._incarnation_id_index = -1
  @@ -117,8 +113,7 @@ class _PerDeviceGenerator(dataset_ops.DatasetV2):
  def _finalize_func(unused_string_handle):
  return array_ops.constant(0, dtypes.int64)
  - finalize_func_concrete = (
  - _finalize_func._get_concrete_function_internal_garbage_collected()) # pylint: disable=protected-access
  + finalize_func_concrete = _finalize_func._get_concrete_function_internal() # pylint: disable=protected-access
  # TODO(b/124254153): Enable autograph once the overhead is low enough.
  @function.defun(
  @@ -131,8 +126,8 @@ class _PerDeviceGenerator(dataset_ops.DatasetV2):
  Tout=[dtypes.int64],
  f=finalize_func_concrete)
  - self._finalize_func = (_remote_finalize_func # pylint: disable=protected-access
  - ._get_concrete_function_internal_garbage_collected())
  + self._finalize_func = (
  + _remote_finalize_func._get_concrete_function_internal()) # pylint: disable=protected-access
  self._finalize_captured_args = self._finalize_func.captured_inputs
  variant_tensor = gen_dataset_ops.generator_dataset(
  