Function get_current_device

src/alignment/model_utils.py:33–35

Get the current device. For GPU we return the local process index to enable multiple GPU training.

()

def get_current_device() -> int:
    """Get the current device. For GPU we return the local process index to enable multiple GPU training."""
    return Accelerator().local_process_index if torch.cuda.is_available() else "cpu"


def get_kbit_device_map() -> Dict[str, int] | None: