Spaces:
Build error
Build error
File size: 1,210 Bytes
583e46a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """
Device detection and management utilities.
"""
import torch
def get_device(prefer_mps: bool = True) -> torch.device:
"""
Auto-detect the best available device.
Priority: CUDA > MPS (Apple Silicon) > CPU
"""
if torch.cuda.is_available():
device = torch.device('cuda')
print(f"Using CUDA: {torch.cuda.get_device_name(0)}")
elif prefer_mps and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = torch.device('mps')
print("Using Apple MPS (Metal Performance Shaders)")
else:
device = torch.device('cpu')
print("Using CPU")
return device
def get_device_info() -> dict:
"""Return device information as a dictionary."""
info = {
'cuda_available': torch.cuda.is_available(),
'mps_available': hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(),
'device': str(get_device()),
}
if torch.cuda.is_available():
info['cuda_device_name'] = torch.cuda.get_device_name(0)
info['cuda_device_count'] = torch.cuda.device_count()
info['cuda_memory_gb'] = round(torch.cuda.get_device_properties(0).total_mem / 1e9, 2)
return info
|