def log_cuda_memory(stage: str = None):
    """
    Log the current CUDA memory usage.

    Args:
        stage (str): The stage of the training process.
    """
    text = "CUDA memory usage"
    if stage is not None:
        text += f" at {stage}"
    log_message(text + ": %.1f GB", torch.cuda.memory_allocated() / GIGABYTE)

    Log the current CUDA memory usage.

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