udpate tests

test/transforms/test_to_device.py

@@ -5,20 +5,18 @@
 
 
 def test_to_device():
-    assert ToDevice(0).__repr__() == 'ToDevice(0)'
+    assert str(ToDevice('cpu')) == 'ToDevice(cpu)'
 
+    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))
-    edge_attr = torch.randn(edge_index.size(1), 8)
 
-    data = Data(edge_index=edge_index, edge_weight=edge_weight,
-                edge_attr=edge_attr, num_nodes=3)
+    data = Data(x=x, edge_index=edge_index, edge_weight=edge_weight)
 
     if (torch.cuda.is_available()):
-        to_device_obj = ToDevice(0)
-        data = to_device_obj(data)
-        assert to_device_obj.device == 0
+        to_device = ToDevice('cuda')
+        data = to_device(data)
+        assert all([val.is_cuda for _, val in data.stores[0].items()])
 
-    to_device_obj = ToDevice('cpu')
-    data = to_device_obj(data)
-    assert to_device_obj.device == 'cpu'
+    to_device = ToDevice('cpu')
+    assert all([val.is_cpu for _, val in data.stores[0].items()])