in local_gemma/utils/config.py [0:0]
def infer_device(device: Optional[str] = None) -> str:
"""
Infers basic devices available on the system. Prioritizes the most performant device.
"""
if device is not None:
return device
elif torch.cuda.is_available():
return "cuda"
elif torch.backends.mps.is_available():
return "mps"
return "cpu"