dgl.multiprocessing.pytorch.call_once_and_shareΒΆ

dgl.multiprocessing.pytorch.call_once_and_share(func, shape, dtype, rank=0)[source]ΒΆ

Invoke the function in a single process of the PyTorch distributed process group, and share the result with other processes.

Parameters

  • func (callable) – Any callable that accepts no arguments and returns an arbitrary object.

  • shape (tuple[int]) – The shape of the shared tensor. Must match the output of func.

  • dtype (torch.dtype) – The data type of the shared tensor. Must match the output of func.

  • rank (int, optional) – The process ID to actually execute the function.