Fix type of HookedTransformerConfig.device

This is either a torch.device or a string like "cpu", but it was
typed as just `Optional[str]`. This fixes it to be
`Optional[Union[str, torch.device]]` and all of the downstream
places that need to be updated.

Found while working on #1219

Author: brendanlong

transformer_lens/HookedTransformerConfig.py

@@ -231,7 +231,7 @@ class HookedTransformerConfig:
     attn_types: Optional[List] = None
     init_mode: str = "gpt2"
     normalization_type: Optional[str] = "LN"
-    device: Optional[str] = None
+    device: Optional[Union[str, torch.device]] = None
     n_devices: int = 1
     attention_dir: str = "causal"
     attn_only: bool = False

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,7 +1015,7 @@ def check_structure(t1: torch.Tensor, t2: torch.Tensor, *, verbose: bool = False
         print(f"column mismatch: {col_mismatch}")
 
 
-def get_device():
+def get_device() -> torch.device:
     if torch.cuda.is_available():
         return torch.device("cuda")
     if torch.backends.mps.is_available() and torch.backends.mps.is_built():
@@ -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