update dataset

shenghuitan · Jan 20, 2021 · 5cb7af9 · 5cb7af9
1 parent 9573d97
commit 5cb7af9
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/tools/utils/static_ps/reader_helper.py b/tools/utils/static_ps/reader_helper.py
@@ -206,11 +206,11 @@ def __init__(self, input_var, file_list, config):
 
     def get_reader(self):
         logger.info("Get Dataset")
-        import paddle.fluid as fluid
-        dataset = fluid.DatasetFactory().create_dataset()
-        dataset.set_use_var(self.input_var)
-        dataset.set_pipe_command(self.pipe_command)
-        dataset.set_batch_size(self.batch_size)
-        dataset.set_thread(self.thread_num)
+        dataset = paddle.distributed.QueueDataset()
+        dataset.init(
+            use_var=self.input_var,
+            pipe_command=self.pipe_command,
+            batch_size=self.batch_size,
+            thread_num=self.thread_num)
         dataset.set_filelist(self.file_list)
         return dataset