graphstorm.utils.setup_device

graphstorm.utils.setup_device(local_rank)

Setup computation device.

Parameters

local_rank: int

Rank of the current process in a distributed environment.

Returns

str: device where the model runs.