Clean up pool properly

        shard_perm = np.random.permutation(num_shards)

      else:

        shard_perm = np.arange(num_shards)

      pool = Pool(1)

      with Pool(1) as pool:

        next_shard = pool.apply_async(dataset.get_shard, (shard_perm[0],))

        for i in range(num_shards):

          X, y, w, ids = next_shard.get()

          if i < num_shards - 1:

          next_shard = pool.apply_async(dataset.get_shard, (shard_perm[i + 1],))

            next_shard = pool.apply_async(dataset.get_shard,

                                          (shard_perm[i + 1],))

          n_samples = X.shape[0]

          # TODO(rbharath): This happens in tests sometimes, but don't understand why?

          # Handle edge case.