baacb20d15066935e4a23c09b1c1a6843331172f,research/object_detection/meta_architectures/center_net_meta_arch.py,,prediction_tensors_to_temporal_offsets,#Any#Any#Any#,338

Before Change

  y_indices = _to_float32(y_indices)
  x_indices = _to_float32(x_indices)

  offsets_flat = _flatten_spatial_dimensions(offset_predictions)

  offsets = tf.gather(offsets_flat, peak_spatial_indices, batch_dims=1)
  return offsets

After Change

  // TF Lite does not support tf.gather with batch_dims > 0, so we need to use
  // tf_gather_nd instead and here we prepare the indices for that.
  combined_indices = tf.stack([
      _multi_range(batch_size, value_repetitions=num_boxes),
      tf.reshape(y_indices, [-1]),
      tf.reshape(x_indices, [-1])
  ], axis=1)

  new_offsets = tf.gather_nd(offset_predictions, combined_indices)
  offsets = tf.reshape(new_offsets, [batch_size, num_boxes, -1])

Instances