# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

from typing import Optional

from lance.dependencies import torch

def preferred_device(device: Optional[str] = None):

"""Get the preferred device for computation.

Parameters

----------

device : str, optional

Device to use for computation. If None, the device will be

detected automatically based on the platform.

Returns

-------

device : torch.device

Device to use for computation.

"""

if device is not None:

if isinstance(device, str):

device = torch.device(device)

return device

if torch.cuda.is_available():

return torch.device("cuda")

if torch.backends.mps.is_available() and torch.backends.mps.is_built():

return torch.device("mps")

return torch.device("cpu")