torch_geometric.transforms.ToDevice

class ToDevice(device: Union[int, str], attrs: Optional[List[str]] = None, non_blocking: bool = False)[source]

Bases: BaseTransform

Performs tensor device conversion, either for all attributes of the Data object or only the ones given by attrs (functional name: to_device).

Parameters

  • device (torch.device) – The destination device.

  • attrs (List[str], optional) – If given, will only perform tensor device conversion for the given attributes. (default: None)

  • non_blocking (bool, optional) – If set to True and tensor values are in pinned memory, the copy will be asynchronous with respect to the host. (default: False)