# wrap_orb.py
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

def wrap_orb_v2(device):
    orbff = pretrained.orb_v2(device=device)
    return ORBCalculator(orbff, device=device)