select_device

select_device(*, on_gpu)[source]

Selects the appropriate device as requested.

Parameters:

on_gpu (bool) – Selects gpu if True.

Returns:

“gpu” if on_gpu is True otherwise returns “cpu”

Return type:

str