dmlc--dgl
Jinjing Zhou
18 行
412 B
Python
18 行
412 B
Python
import dgl
|
|
import backend as F
|
|
import unittest
|
|
|
|
|
|
@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
|
|
def test_to_device():
|
|
g = dgl.DGLGraph()
|
|
g.add_nodes(5, {'h' : F.ones((5, 2))})
|
|
g.add_edges([0, 1], [1, 2], {'m' : F.ones((2, 2))})
|
|
if F.is_cuda_available():
|
|
g = g.to(F.cuda())
|
|
assert g is not None
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_to_device()
|