Method fetch

python/dgl/storages/pytorch_tensor.py:31–58

(self, indices, device, pin_memory=False, **kwargs)

    """Feature storages for slicing a PyTorch tensor."""

    def fetch(self, indices, device, pin_memory=False, **kwargs):
        device = torch.device(device)
        storage_device_type = self.storage.device.type
        indices_device_type = indices.device.type
        if storage_device_type != "cuda":
            if indices_device_type == "cuda":
                if self.storage.is_pinned():
                    return gather_pinned_tensor_rows(self.storage, indices)
                else:
                    raise ValueError(
                        f"Got indices on device {indices.device} whereas the feature tensor "
                        f"is on {self.storage.device}. Please either (1) move the graph "
                        f"to GPU with to() method, or (2) pin the graph with "
                        f"pin_memory_() method."
                    )
            # CPU to CPU or CUDA - use pin_memory and async transfer if possible
            else:
                return _fetch_cpu(
                    indices,
                    self.storage,
                    self.storage.shape[1:],
                    device,
                    pin_memory,
                    **kwargs,
                )
        else:
            # CUDA to CUDA or CPU
            return _fetch_cuda(indices, self.storage, device, **kwargs)