onf-c0/src/onf/device.py

"""The one torch.device resolution helper for the whole package.

Lives at the top level rather than under onf.graph because onf.field needs it too, and
onf.graph already depends on onf.field -- putting it in the graph package would make that a
cycle. Imports torch and nothing from onf.
"""

from __future__ import annotations

import torch


def resolve_device(device: str | torch.device | None) -> torch.device:
    """None -> CUDA if available else CPU; a string/torch.device is used as given."""
    if isinstance(device, torch.device):
        return device
    if device:
        return torch.device(device)
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")