INV / helium /device_utils.py
Fred808's picture
Upload 256 files
7a0c684 verified
from __future__ import annotations
from typing import Optional
from .virtual_gpu_device import VirtualGPUDevice
_default_device = None
def get_default_device() -> VirtualGPUDevice:
"""Get the default virtual GPU device"""
global _default_device
if _default_device is None:
_default_device = VirtualGPUDevice()
return _default_device
def get_gpu_device(device_id: Optional[int] = None) -> VirtualGPUDevice:
"""Get a specific virtual GPU device by ID"""
if device_id is None:
return get_default_device()
return VirtualGPUDevice(device_id)