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 byattrs
(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
)