src/transformers/models/esm/openfold_utils/chunk_utils.py/_fetch_dims
def _fetch_dims(tree: Union[dict, list, tuple, torch.Tensor]) -> List[Tuple[int, ...]]:
    shapes = []
    if isinstance(tree, dict):
        for v in tree.values():
            shapes.extend(_fetch_dims(v))
    elif isinstance(tree, (list, tuple)):
        for t in tree:
            shapes.extend(_fetch_dims(t))
    elif isinstance(tree, torch.Tensor):
        shapes.append(tree.shape)
    else:
        raise ValueError("Not supported")

    return shapes