As part of torch 1.9.0 we are introducing torch.distributed.run to replace torch.distributed.launch definition is here (pytorch/run.py at master · pytorch/pytorch · GitHub 16). This eventually calls into a function called elastic_launch (pytorch/api.py at master · pytorch/pytorch · GitHub 7) which seems to be what you are looking for. For example, you can import it and use it like elastic_launch(LaunchConfig, your_module_function)(your_module_args). Does this satisfy your use case?