Fix type of HookedTransformerConfig.device

This is typed as `Optional[str]` but sometimes returns `torch.device`.
Updated the code to just return the `str` instead of wrapping with a
device.

I'm not confident that every function which takes a device will
always be passed a string, so I didn't change functions like
warn_if_mps.

Found while working on #1219

Author: brendanlong

tests/unit/pretrained_weight_conversions/test_apertus.py

@@ -183,7 +183,7 @@ def test_zero_biases_have_correct_device(self):
             "blocks.0.mlp.b_out",
             "unembed.b_U",
         ]:
-            assert sd[key].device.type == cfg.device.type, f"{key} on wrong device"
+            assert sd[key].device.type == cfg.device, f"{key} on wrong device"
 
     def test_unembed_shapes(self):
         cfg = make_cfg()

transformer_lens/train.py

@@ -5,7 +5,7 @@
 """
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 import torch.optim as optim
@@ -50,7 +50,7 @@ class HookedTransformerTrainConfig:
     max_grad_norm: Optional[float] = None
     weight_decay: Optional[float] = None
     optimizer_name: str = "Adam"
-    device: Optional[str] = None
+    device: Optional[Union[str, torch.device]] = None
     warmup_steps: int = 0
     save_every: Optional[int] = None
     save_dir: Optional[str] = None

transformer_lens/utils.py

@@ -1015,9 +1015,9 @@ def check_structure(t1: torch.Tensor, t2: torch.Tensor, *, verbose: bool = False
         print(f"column mismatch: {col_mismatch}")
 
 
-def get_device():
+def get_device() -> str:
     if torch.cuda.is_available():
-        return torch.device("cuda")
+        return "cuda"
     if torch.backends.mps.is_available() and torch.backends.mps.is_built():
         major_version = int(torch.__version__.split(".")[0])
         if major_version >= 2:
@@ -1026,17 +1026,17 @@ def get_device():
                 _MPS_MIN_SAFE_TORCH_VERSION is not None
                 and _torch_version_tuple() >= _MPS_MIN_SAFE_TORCH_VERSION
             ):
-                return torch.device("mps")
+                return "mps"
             if os.environ.get("TRANSFORMERLENS_ALLOW_MPS", "") == "1":
-                return torch.device("mps")
+                return "mps"
             logging.info(
                 "MPS device available but not auto-selected due to known correctness issues "
                 "(PyTorch %s). Set TRANSFORMERLENS_ALLOW_MPS=1 to override. See: "
                 "https://github.com/TransformerLensOrg/TransformerLens/issues/1178",
                 torch.__version__,
             )
 
-    return torch.device("cpu")
+    return "cpu"
 
 
 _mps_warned = False
@@ -1051,7 +1051,7 @@ def _torch_version_tuple() -> tuple[int, ...]:
     return tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:2])
 
 
-def warn_if_mps(device):
+def warn_if_mps(device: Union[str, torch.device]) -> None:
     """Emit a one-time warning if device is MPS and TRANSFORMERLENS_ALLOW_MPS is not set.
 
     Automatically suppressed when the installed PyTorch version meets or exceeds