import torch
def get_best_device(prefer_cpu=False, allow_mps=True):
if prefer_cpu:
return torch.device("cpu")
if torch.cuda.is_available():
return torch.device("cuda")
if allow_mps and torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")
def get_device_info():
cuda_available = torch.cuda.is_available()
mps_available = torch.backends.mps.is_available()
if cuda_available:
best = "cuda"
elif mps_available:
best = "mps"
else:
best = "cpu"
return {
"cuda_available": cuda_available,
"mps_available": mps_available,
"best_device": best,
}
def log_device_info(logger):
info = get_device_info()
if info["cuda_available"]:
logger.info("CUDA GPU available")
elif info["mps_available"]:
logger.info("MPS (Apple GPU) available")
else:
logger.info("No GPU available - will use CPU")