import tensorflow as tf

      dataset = dataset.map(map_fn, num_parallel_calls=16)

    else:
      types = {movielens.USER_COLUMN: rconst.USER_DTYPE,
               movielens.ITEM_COLUMN: rconst.ITEM_DTYPE}
      shapes = {movielens.USER_COLUMN: tf.TensorShape([batch_size]),
                movielens.ITEM_COLUMN: tf.TensorShape([batch_size])}

      if self._is_training:
        types[rconst.VALID_POINT_MASK] = np.bool
        shapes[rconst.VALID_POINT_MASK] = tf.TensorShape([batch_size])

        types = (types, np.bool)
        shapes = (shapes, tf.TensorShape([batch_size]))

      else:
        types[rconst.DUPLICATE_MASK] = np.bool
        shapes[rconst.DUPLICATE_MASK] = tf.TensorShape([batch_size])

      data_generator = functools.partial(
