Function is_on_gpu

bitsandbytes/functional.py:351–390

Verifies that the input tensors are all on the same device. An input tensor may also be marked as `paged`, in which case the device placement is ignored. CPU tensors are allowed and checked for consistency among themselves. Args: tensors (`Iterable[Optional[torch.Tensor]]`): A

(tensors: Iterable[Optional[torch.Tensor]])

def is_on_gpu(tensors: Iterable[Optional[torch.Tensor]]):
    """Verifies that the input tensors are all on the same device.

    An input tensor may also be marked as `paged`, in which case the device placement is ignored.
    CPU tensors are allowed and checked for consistency among themselves.

    Args:
        tensors (`Iterable[Optional[torch.Tensor]]`): A list of tensors to verify.

    Raises:
        `RuntimeError`: Raised when the verification fails.

    Returns:
        `Literal[True]`
    """

    devices = set()

    for t in tensors:
        # NULL pointers and paged tensors are OK.
        if t is not None and not getattr(t, "is_paged", False):
            devices.add((t.device.type, t.device.index))

    # All tensors on CPU is valid
    if devices == {("cpu", None)}:
        return True

    # Check that no CPU tensors are mixed with GPU tensors
    has_cpu = ("cpu", None) in devices
    if has_cpu and len(devices) > 1:
        raise RuntimeError(
            f"Input tensors need to be on the same device, but found the following tensor and device combinations:\n {[(t.shape, t.device) for t in tensors if t is not None]}",
        )

    # GPU path: all tensors must be on the same single GPU
    if len(devices) > 1:
        raise RuntimeError(
            f"Input tensors need to be on the same GPU, but found the following tensor and device combinations:\n {[(t.shape, t.device) for t in tensors if t is not None]}",
        )
    return True


def _get_tensor_stream(tensor: Tensor) -> ct.c_void_p: