def rearrange(rearrange_order, dim=0, *inputs):
    """Rearrages input tensors based on an order and along a given dimension.

  Args:
    rearrange_order: Order to use while rearranging.
    dim: Dimension along which to rearrange.
    *inputs: List of input tensors.
  """
    rearranged_inputs = []
    for input_tensor in inputs:
        assert (
            input_tensor.shape[dim] == rearrange_order.shape[0]
        ), "Rearrange " "along dim {0} is incompatible!".format(dim)
        rearranged_inputs.append(input_tensor.index_select(dim, rearrange_order))
    return tuple(rearranged_inputs)