torch_geometric.transforms.ToDevice

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

Bases: BaseTransform

执行张量设备转换，可以针对Data对象的所有属性，或者仅针对由attrs指定的属性（功能名称：to_device）。

Parameters:

  • device (torch.device) – 目标设备。

  • attrs (List[str], optional) – 如果提供，将仅对给定的属性执行张量设备转换。（默认值：None）

  • non_blocking (bool, 可选) – 如果设置为 True 并且张量 值在固定内存中，复制将与主机异步进行。（默认值：False）