diff --git a/PaddleNLP/examples/text_graph/erniesage/config/erniesage_link_prediction.yaml b/PaddleNLP/examples/text_graph/erniesage/config/erniesage_link_prediction.yaml
index bde00ad24dcb60019635f3b8e4fce3b06bf2c2b9..3beb069111b093e20e81a87677eec511aeff7455 100644
--- a/PaddleNLP/examples/text_graph/erniesage/config/erniesage_link_prediction.yaml
+++ b/PaddleNLP/examples/text_graph/erniesage/config/erniesage_link_prediction.yaml
@@ -1,7 +1,7 @@
 # Global Enviroment Settings 
 
 # trainer config ------
-n_gpu: 2 # delete it, if use cpu to train
+n_gpu: 1 # number of gpus used to train, delete it, if use cpu
 seed: 2020
 
 task: "link_prediction"
diff --git a/PaddleNLP/examples/text_graph/erniesage/data/dataset.py b/PaddleNLP/examples/text_graph/erniesage/data/dataset.py
index 173af3dbb1071a733df71ed226a2e61ce9f7acc6..91eeed1af90303552be7c476bc4381d06399f3f4 100644
--- a/PaddleNLP/examples/text_graph/erniesage/data/dataset.py
+++ b/PaddleNLP/examples/text_graph/erniesage/data/dataset.py
@@ -21,7 +21,7 @@ from paddle.io import Dataset, IterableDataset
 from paddlenlp.utils.log import logger
 import pgl
 from pgl import Graph
-from pgl.nn.sampling import graphsage_sample
+from pgl.sampling import graphsage_sample
 
 __all__ = [
     "TrainData",