torchfoo.current_device

torchfoo.current_device()

Return the current device.

In a distributed context, returns the device assigned to the current process. Outside of a distributed context, returns the current CUDA device if CUDA is available, otherwise torch.device("cpu").

Returns:

The current torch.device, or an int CUDA device index when running in a distributed CUDA context.