dgl.multiprocessing.pytorch.shared_tensorΒΆ

dgl.multiprocessing.pytorch.shared_tensor(shape, dtype=torch.float32)[source]ΒΆ

Create a tensor in shared memory accessible by all processes within the same torch.distributed process group.

The content is uninitialized.

Parameters

  • shape (tuple[int]) – The shape of the tensor.

  • dtype (torch.dtype, optional) – The dtype of the tensor.

Returns

The shared tensor.

Return type

Tensor