refactor

deepmodeling · Dec 19, 2024 · 96c9f03 · 96c9f03
1 parent 5ad15d1
commit 96c9f03
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py
@@ -101,10 +101,9 @@ def construct_dataset(system):
         if global_rank == 0:
             with Pool(os.cpu_count()) as pool:
                 self.systems = pool.map(construct_dataset, systems)
-            if dist.is_initialized():
-                dist.broadcast_object_list(self.systems)
         else:
             self.systems = [None] * len(systems) # type: ignore
+        if dist.is_initialized():
             dist.broadcast_object_list(self.systems)
             assert self.systems[-1] is not None
         self.sampler_list: list[DistributedSampler] = []