def METHOD_NAME(self, x: Dict, indices: List):
    """
    Sort the values of each key according to the indices.
    This is to make sure that prediction results follow the order of input samples.
    Parameters
    ----------
    x
        A dictionary with all the predictions.
    indices
        A list of indices.
    """
    results = dict()
    for k, v in x.items():
        if isinstance(v, dict):
            results[k] = self.METHOD_NAME(v, indices)
        else:
            assert len(indices) == len(v)
            results[k] = [x for _, x in sorted(zip(indices, v), key=lambda ele: ele[0])]
            results[k] = torch.stack(results[k])
    return results