def fn(i):
    num_valid_voxels_i = num_valid_voxels[i]
    num_valid_points_i = num_valid_points[i]
    voxel_features_i = voxel_features[i, :num_valid_voxels_i, :]
    segment_ids_i = segment_ids[i, :num_valid_points_i]
    point_features = tf.gather(voxel_features_i, segment_ids_i)
    point_features_rank = len(point_features.get_shape().as_list())
    point_features_paddings = [[0, num_points - num_valid_points_i]]
    for _ in range(point_features_rank - 1):
      point_features_paddings.append([0, 0])
    point_features = tf.pad(point_features, paddings=point_features_paddings)
    return point_features