test_to_device.py

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