def log_cuda_max_memory(stage: str = None):
    """
    Log the max CUDA memory usage.

    Args:
        stage (str): The stage of the training process.
    """
    torch.cuda.synchronize()
    max_memory_allocated = torch.cuda.max_memory_allocated()
    max_memory_reserved = torch.cuda.max_memory_reserved()
    log_message("CUDA max memory max memory allocated at " + stage + ": %.1f GB", max_memory_allocated / GIGABYTE)
    log_message("CUDA max memory max memory reserved at " + stage + ": %.1f GB", max_memory_reserved / GIGABYTE)

    Log the max CUDA memory usage.

    Args:
        stage (str): The stage of the training process.
    
opensora/utils/misc.py