def rank_zero_only(fn: Callable[..., R]) -> Callable[..., Optional[R]]:
    """Decorator to ensure the function only runs on the process with rank 0.

    Example:
        >>> @rank_zero_only
        >>> def print_rank_zero(message: str):
        >>>     print(message)
        >>>
        >>> print_rank_zero("Hello from rank 0!")
    """

    def wrapped(*args: Any, **kwargs: Any) -> Optional[R]:
        if rank() == 0:
            return fn(*args, **kwargs)
        return None

    return wrapped
Decorator to ensure the function only runs on the process with rank 0.

    Example:
        >>> @rank_zero_only
        >>> def print_rank_zero(message: str):
        >>>     print(message)
        >>>
        >>> print_rank_zero("Hello from rank 0!")
    
lightly/utils/dist.py