Skip to content

declearn.utils.set_device_policy

Update the current global device policy.

To access the current policy, use declearn.utils.set_device_policy.

Parameters:

Name Type Description Default
gpu bool

Whether to use a GPU device rather than the CPU one to back data and computations. If no GPU is available, use CPU with a warning.

 required
idx Optional[int]

Optional index of the GPU device to use. If this index exceeds the number of available GPUs, select one arbitrarily, with a warning.

 None
Source code in declearn/utils/_device_policy.py 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def set_device_policy(
    gpu: bool,
    idx: Optional[int] = None,
) -> None:
    """Update the current global device policy.

    To access the current policy, use `declearn.utils.set_device_policy`.

    Parameters
    ----------
    gpu: bool
        Whether to use a GPU device rather than the CPU one to back data
        and computations. If no GPU is available, use CPU with a warning.
    idx: int or None, default=None
        Optional index of the GPU device to use.
        If this index exceeds the number of available GPUs, select one
        arbitrarily, with a warning.
    """
    # Using a global statement to have a proper setter to a private variable.
    global DEVICE_POLICY  # pylint: disable=global-statement
    DEVICE_POLICY = DevicePolicy(gpu, idx)