gigl.distributed.utils.device

Functions

get_available_device(local_process_rank)	Returns the available device for the current process.
get_device_from_process_group()	Returns the device for the current process group.

Module Contents

gigl.distributed.utils.device.get_available_device(local_process_rank)

Returns the available device for the current process.

Parameters:

local_process_rank (int) – The local rank of the current process within a node.

Returns:

The device to use.

Return type:

torch.device

gigl.distributed.utils.device.get_device_from_process_group()

Returns the device for the current process group. :param None:

Raises:

ValueError – If the distributed environment is not initialized.

Returns:

The device to use.

Return type:

torch.device