1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
|
import torch
from torch_geometric.data import Data
from torch_geometric.testing import withDevice
from torch_geometric.transforms import ToDevice
@withDevice
def test_to_device(device):
x = torch.randn(3, 4)
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
edge_weight = torch.randn(edge_index.size(1))
data = Data(x=x, edge_index=edge_index, edge_weight=edge_weight)
transform = ToDevice(device)
assert str(transform) == f'ToDevice({device})'
data = transform(data)
for key, value in data:
assert value.device == device
|
