Match device selection of TransformerBridge to HookedTransformer (#1047)

Co-authored-by: Bryce Meyer <[email protected]>

Author: degenfabian

transformer_lens/model_bridge/sources/transformers.py

@@ -18,7 +18,7 @@
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge.bridge import TransformerBridge
 from transformer_lens.supported_models import MODEL_ALIASES
-from transformer_lens.utils import get_tokenizer_with_bos
+from transformer_lens.utils import get_device, get_tokenizer_with_bos
 
 
 def map_default_transformer_lens_config(hf_config):
@@ -210,9 +210,12 @@ def boot(
 
     adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config)
 
+    # No device specified by user, use the best available device for the current system
+    if device is None:
+        device = get_device()
+
     # Add device information to the config
-    if device is not None:
-        adapter.cfg.device = str(device)
+    adapter.cfg.device = str(device)
 
     # Load the model from HuggingFace using the original config
     hf_model = AutoModelForCausalLM.from_pretrained(
@@ -221,9 +224,8 @@ def boot(
         torch_dtype=dtype,
     )
 
-    # Move model to device if specified
-    if device is not None:
-        hf_model = hf_model.to(device)
+    # Move model to device
+    hf_model = hf_model.to(device)
 
     # Load the tokenizer
     tokenizer = tokenizer