diff --git a/python/dglke/dataloader/sampler.py b/python/dglke/dataloader/sampler.py
index 38ff203..2e46bfd 100644
--- a/python/dglke/dataloader/sampler.py
+++ b/python/dglke/dataloader/sampler.py
@@ -335,7 +335,7 @@ def ConstructGraph(dataset, args):
     n_entities = dataset.n_entities
 
     coo = sp.sparse.coo_matrix((np.ones(len(src)), (src, dst)), shape=[n_entities, n_entities])
-    g = dgl.DGLGraph(coo, readonly=True, multigraph=True, sort_csr=True)
+    g = dgl._deprecate.graph.DGLGraph(coo, readonly=True, multigraph=True, sort_csr=True)
     g.edata['tid'] = F.tensor(etype_id, F.int64)
     if args.has_edge_importance:
         e_impts = F.tensor(dataset.train[3], F.float32)
@@ -418,7 +418,7 @@ def create_sampler(self, batch_size, neg_sample_size=2, neg_chunk_size=None, mod
                            exclude_positive=exclude_positive,
                            return_false_neg=False)
 
-class ChunkNegEdgeSubgraph(dgl.DGLGraph):
+class ChunkNegEdgeSubgraph(dgl._deprecate.graph.DGLGraph):
     """Wrapper for negative graph
 
         Parameters
diff --git a/python/dglke/models/ke_model.py b/python/dglke/models/ke_model.py
index f17f104..f01000f 100644
--- a/python/dglke/models/ke_model.py
+++ b/python/dglke/models/ke_model.py
@@ -84,7 +84,7 @@ def attach_graph(self, g, etid_field='tid', ntid_filed='ntid'):
         """
         self._etid_field = etid_field
         self._ntid_filed = ntid_filed
-        assert isinstance(g, dgl.DGLGraph)
+        assert isinstance(g, dgl._deprecate.graph.DGLGraph)
         self._g = g
 
     def load(self, model_path):
diff --git a/python/dglke/partition.py b/python/dglke/partition.py
index 9743d83..0ed15ed 100644
--- a/python/dglke/partition.py
+++ b/python/dglke/partition.py
@@ -111,7 +111,7 @@ def main():
     src, etype_id, dst = dataset.train
     coo = sp.sparse.coo_matrix((np.ones(len(src)), (src, dst)),
             shape=[dataset.n_entities, dataset.n_entities])
-    g = dgl.DGLGraph(coo, readonly=True, multigraph=True, sort_csr=True)
+    g = dgl._deprecate.graph.DGLGraph(coo, readonly=True, multigraph=True, sort_csr=True)
     g.edata['tid'] = F.tensor(etype_id, F.int64)
 
     print('partition graph...')
diff --git a/python/dglke/tests/test_score.py b/python/dglke/tests/test_score.py
index 5915e36..1c16fe6 100644
--- a/python/dglke/tests/test_score.py
+++ b/python/dglke/tests/test_score.py
@@ -50,7 +50,7 @@ class dotdict(dict):
 
 def generate_rand_graph(n, func_name):
     arr = (sp.sparse.random(n, n, density=0.1, format='coo') != 0).astype(np.int64)
-    g = dgl.DGLGraph(arr, readonly=True)
+    g = dgl._deprecate.graph.DGLGraph(arr, readonly=True)
     num_rels = 10
     entity_emb = F.uniform((g.number_of_nodes(), 10), F.float32, F.cpu(), 0, 1)
     if func_name == 'RotatE':
diff --git a/python/dglke/tests/test_topk.py b/python/dglke/tests/test_topk.py
index 945bfc3..986612f 100644
--- a/python/dglke/tests/test_topk.py
+++ b/python/dglke/tests/test_topk.py
@@ -562,7 +562,7 @@ def check_topk_score2(score_model, exclude_mode):
                     th.full((num_entity*2,), 2, dtype=th.long),
                     th.full((num_entity*2,), 3, dtype=th.long)],
                     dim=0)
-    g = dgl.graph((src, dst))
+    g = dgl._deprecate.graph.DGLGraph((src, dst))
     g.edata['tid'] = etype
 
     _check_topk_score2(score_model, g, num_entity, num_rels, exclude_mode)
@@ -1261,4 +1261,4 @@ def test_gnn_model_topk_emb(device='cpu'):
     test_rescal_model_topk_emb(device=0)
     #test_transr_model_topk_emb()
     test_rotate_model_topk_emb(device=0)
-    test_gnn_model_topk_emb(device=0)
\ No newline at end of file
+    test_gnn_model_topk_emb(device=0)
diff --git a/tests/scripts/task_kg_test.sh b/tests/scripts/task_kg_test.sh
index d4f702a..471d34e 100644
--- a/tests/scripts/task_kg_test.sh
+++ b/tests/scripts/task_kg_test.sh
@@ -34,10 +34,10 @@ conda activate ${DGLBACKEND}-ci
 # test
 if [ "$2" == "cpu" ]; then
     pip3 uninstall dgl
-    pip3 install dgl==0.4.3post2
+    pip3 install dgl>=0.5.0,<=0.9.1
 else
     pip3 uninstall dgl
-    pip3 install dgl-cu102==0.4.3post2
+    pip3 install dgl-cu102>=0.5.0,<=0.9.1
 fi
 
 pushd $KG_DIR> /dev/null