def get_gpu_memory(device: str | torch.device = "cuda", unit: str = "GiB") -> int:
    """Get the GPU memory of the current device.

    Args:
        device (`str` | `torch.device`, optional, defaults to `"cuda"`):
            device.

    Returns:
        `int`:
            GPU memory in bytes.
    """
    if isinstance(device, str):
        device = torch.device(device)
    assert unit in ("GiB", "MiB", "B")
    memory = torch.cuda.get_device_properties(device).total_memory
    if unit == "GiB":
        return memory // (1024**3)
    elif unit == "MiB":
        return memory // (1024**2)
    else:
        return memory
Get the GPU memory of the current device.

    Args:
        device (`str` | `torch.device`, optional, defaults to `"cuda"`):
            device.

    Returns:
        `int`:
            GPU memory in bytes.
    
nunchaku/utils.py