fix shard_dataloader (#65262)

zhiqiu · web-flow · commit 3dd4cbfb7989 · 2024-06-19T08:50:57.000+08:00
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -2719,7 +2719,6 @@ def __len__(self):
         return len(self._dataloader)
 
     def __iter__(self):
-        self.iter = self._dataloader.__iter__()
         return self
 
     def _get_mesh_and_placement(self, index):
@@ -2840,7 +2839,8 @@ def __next__(self):
         return self._get_batch(batch_data)
 
     def __call__(self):
-        return self.__iter__()
+        self.iter = self._dataloader.__iter__()
+        return self
 
 
 def shard_dataloader(
diff --git a/python/paddle/distributed/auto_parallel/static/engine.py b/python/paddle/distributed/auto_parallel/static/engine.py
@@ -283,7 +283,7 @@ def __init__(
     def _prepare_data_spec_from_dataloader(self, dataloader):
         inputs_spec = []
         labels_spec = []
-        data = next(iter(dataloader))
+        data = next(dataloader())
         if hasattr(dataloader, "batch_sampler"):
             batch_sampler = dataloader.batch_sampler
         else:
diff --git a/test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model.py b/test/auto_parallel/hybrid_strategy/test_semi_auto_parallel_llama_model.py
@@ -217,8 +217,6 @@ def setUp(self):
         }
 
     def test_simple_net_hybrid_strategy(self):
-        # Temporally disable because of random failure
-        return
         envs_list = test_base.gen_product_envs_list(
             self._default_envs, self._changeable_envs
         )

Original file line number	Diff line number	Diff line change
`@@ -217,8 +217,6 @@ def setUp(self):`
`217`	`217`	`}`
`218`	`218`
`219`	`219`	`def test_simple_net_hybrid_strategy(self):`
`220`		`- # Temporally disable because of random failure`
`221`		`- return`
`222`	`220`	`envs_list = test_base.gen_product_envs_list(`
`223`	`221`	`self._default_envs, self._changeable_envs`
`224`	`222`	`)`