def test__hvd_dist_model_spawn_cuda():
num_workers_per_machine = torch.cuda.device_count()
_HorovodDistModel.spawn(
_test_dist_spawn_fn,
args=("horovod", num_workers_per_machine, "cuda"),
kwargs_dict={},
nproc_per_node=num_workers_per_machine,
use_gloo=True,
)
def test__hvd_dist_model_spawn():
num_workers_per_machine = 4
_HorovodDistModel.spawn(
_test_dist_spawn_fn,
args=("horovod", num_workers_per_machine, "cpu"),
kwargs_dict={},
nproc_per_node=num_workers_per_machine,
use_gloo=True,
)