def helper_sf_gather(sf_tensor, indeces):
    """
    Helper function, given a SF tensor and action indeces, return only the
    SFs associated with the actions
    :param sf_tensor: SF tensor with action specific dimension of -2, here
                      with size (batch_n, *, |A|, d)
    :param indeces: action indeces, size (batch_n, *, 1)
    :return: action-sf tensor of size (batch_n, *, d)
    """
    sizes = list(sf_tensor.size())  # list: [batch, *, |A|, d]
    sizes[-2] = 1  # list: [batch, *, 1, d]

    idxs = indeces.clone().unsqueeze(-1)  # (batch_n, 1, 1)
    while len(sizes) > len(idxs.size()):
        idxs = idxs.unsqueeze(-2)  # (batch_n, *(1), 1, 1)
    idxs = idxs.expand(sizes)  # (batch_n, *, 1, d)

    # gathering. For dim == -2
    # out[i]...[j][k] = input[i]...[idxs[i]...[j][k]][k]
    sf_a_tensor = sf_tensor.gather(-2, idxs)  # (batch_n, *, 1, d)
    return sf_a_tensor.squeeze(-2)