def test_to_device(index_dtype):
g1 = dgl.heterograph({('user', 'plays', 'game'): [(0, 0), (1, 1)]},
index_dtype=index_dtype)
g1.nodes['user'].data['h1'] = F.copy_to(F.tensor([[0.], [1.]]), F.cpu())
g1.nodes['user'].data['h2'] = F.copy_to(F.tensor([[3.], [4.]]), F.cpu())
g1.edges['plays'].data['h1'] = F.copy_to(F.tensor([[2.], [3.]]), F.cpu())
g2 = dgl.heterograph({('user', 'plays', 'game'): [(0, 0), (1, 0)]},
index_dtype=index_dtype)
g2.nodes['user'].data['h1'] = F.copy_to(F.tensor([[1.], [2.]]), F.cpu())
g2.nodes['user'].data['h2'] = F.copy_to(F.tensor([[4.], [5.]]), F.cpu())
g2.edges['plays'].data['h1'] = F.copy_to(F.tensor([[0.], [1.]]), F.cpu())
bg = dgl.batch_hetero([g1, g2])
if F.is_cuda_available():
bg1 = bg.to(F.cuda())
assert bg1 is not None
assert bg.batch_size == bg1.batch_size
assert bg.batch_num_nodes('user') == bg1.batch_num_nodes('user')
assert bg.batch_num_edges('plays') == bg1.batch_num_edges('plays')
# set feature
g1 = dgl.heterograph({('user', 'plays', 'game'): [(0, 0), (1, 1)]},
index_dtype=index_dtype)
g2 = dgl.heterograph({('user', 'plays', 'game'): [(0, 0), (1, 0)]},
index_dtype=index_dtype)
bg = dgl.batch_hetero([g1, g2])
if F.is_cuda_available():
bg1 = bg.to(F.cuda())
bg1.nodes['user'].data['test'] = F.copy_to(F.tensor([0, 1, 2, 3]),
F.cuda())
bg1.edata['test'] = F.copy_to(F.tensor([0, 1, 2, 3]), F.cuda())
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
def test_to_device():
hg = create_test_heterograph()
if F.is_cuda_available():
hg = hg.to(F.cuda())
assert hg is not None