Fix (#1684)

tests/compute/test_batched_heterograph.py

 import dgl
 import backend as F
+import unittest
 
 from dgl.base import ALL
 from utils import parametrize_dtype
@@ -268,8 +269,35 @@ def test_batching_with_zero_nodes_edges(index_dtype):
     g2.nodes['u'].data['x'] = F.tensor([1])
     dgl.batch_hetero([g1, g2])
 
+@unittest.skipIf(F._default_context_str == 'cpu', reason="Need gpu for this test")
+@parametrize_dtype
+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')
+
 if __name__ == '__main__':
     test_batching_hetero_topology()
     test_batching_hetero_and_batched_hetero_topology()
     test_batched_features()
     test_batching_with_zero_nodes_edges()
+    # test_to_device()