utils.validation.set_cpu_gpu_device

utils.validation.set_cpu_gpu_device(estimator, device='cpu')

Set the device for the estimator to either ‘cpu’, ‘gpu’, ‘cuda’, or None.

Parameters

Name Type Description Default
estimator object Estimator compatible with the scikit-learn API. required
device str | None Device to set. Options are ‘cpu’, ‘gpu’, ‘cuda’, or None. Defaults to ‘cpu’. 'cpu'

Returns

Name Type Description
str | None The device that was set on the estimator before the function was called.