fix: use device_map for HF model loading (#581) (#587)

planetf1 · web-flow · commit 8a385d50e153 · 2026-03-06T17:16:57.000Z
diff --git a/mellea/backends/huggingface.py b/mellea/backends/huggingface.py
@@ -270,8 +270,8 @@ def __init__(
                 )
                 # Get the model and tokenizer.
                 self._model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
-                    self._hf_model_id
-                ).to(self._device)  # type: ignore
+                    self._hf_model_id, device_map=str(self._device)
+                )
                 self._tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(
                     self._hf_model_id
                 )

Original file line number	Diff line number	Diff line change
`@@ -270,8 +270,8 @@ def __init__(`
`270`	`270`	`)`
`271`	`271`	`# Get the model and tokenizer.`
`272`	`272`	`self._model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(`
`273`		`- self._hf_model_id`
`274`		`- ).to(self._device) # type: ignore`
	`273`	`+ self._hf_model_id, device_map=str(self._device)`
	`274`	`+ )`
`275`	`275`	`self._tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(`
`276`	`276`	`self._hf_model_id`
`277`	`277`	`)`