fix: take cfg.device from a pre-loaded hf_model instead of get_device()

tests/unit/model_bridge/test_boot_device.py

@@ -0,0 +1,22 @@
+"""boot()'s device for a pre-loaded ``hf_model``."""
+
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from transformer_lens.model_bridge import TransformerBridge
+
+
+def test_device_defaults_to_hf_model_device(monkeypatch):
+    # boot never moves hf_model, so cfg.device must be the model's device, not get_device()'s.
+    monkeypatch.setattr(
+        "transformer_lens.model_bridge.sources.transformers.source.get_device", lambda: "cuda"
+    )
+    hf_config = AutoConfig.for_model(
+        "gpt_neox", hidden_size=32, num_hidden_layers=1, num_attention_heads=4, intermediate_size=64
+    )
+    with torch.device("meta"):
+        hf_model = AutoModelForCausalLM.from_config(hf_config)
+
+    bridge = TransformerBridge.boot_transformers("EleutherAI/pythia-70m", hf_model=hf_model)
+
+    assert bridge.cfg.device == "meta"

transformer_lens/model_bridge/sources/transformers/source.py

@@ -259,9 +259,10 @@ def boot(
         n_devices, device_map, device, max_memory
     )
     if resolved_device_map is None:
-        if device is None:
+        # hf_model is never moved; without device= cfg.device is read from it below.
+        if device is None and hf_model is None:
             device = get_device()
-        adapter.cfg.device = str(device)
+        adapter.cfg.device = None if device is None else str(device)
     else:
         # cfg.device set from hf_device_map after the model is loaded; provisionally None.
         adapter.cfg.device = None