def METHOD_NAME(self, device_type: str, device_id: int):
    if device_type == "cpu":
        self._model = self.copy(mx.cpu())
    elif device_type == "gpu":
        self._model = self.copy(mx.gpu())
    else:
        msg = f"Unexpected device_type: {device_type}. Try 'cpu' or 'gpu'."
        raise ValueError(msg)