diff --git a/dygraph/train.py b/dygraph/train.py
index a776cb0b5c0f3e170835633a58637af26b499bee..17ba2d68730e39c480c084c0de1afc790582ba41 100644
--- a/dygraph/train.py
+++ b/dygraph/train.py
@@ -21,7 +21,7 @@ import paddle.fluid as fluid
 from paddle.fluid.dygraph.parallel import ParallelEnv
 from paddle.fluid.io import DataLoader
 
-from datasets import OpticDiscSeg, Dataset
+from datasets import OpticDiscSeg
 import transforms as T
 import models
 import utils.logging as logging