Corrected work dict for MPI tasks.

这个提交包含在:
Craig Warren
2022-10-25 14:27:07 +01:00
父节点 9d840ca714
当前提交 1a19459325

查看文件

@@ -109,11 +109,17 @@ class MPIContext(Context):
self.rank = self.comm.rank self.rank = self.comm.rank
self.MPIExecutor = MPIExecutor self.MPIExecutor = MPIExecutor
def _run_model(self, i): def _run_model(self, **work):
"""Process for running a single model.""" """Process for running a single model.
Args:
work (dict): contains any additional information that is passed to
MPI workers. By default only model number (i) is
used.
"""
# Create configuration for model # Create configuration for model
config.model_num = i config.model_num = work['i']
model_config = config.ModelConfig() model_config = config.ModelConfig()
# Set GPU deviceID according to worker rank # Set GPU deviceID according to worker rank
if config.sim_config.general['cuda']: if config.sim_config.general['cuda']: