Remove n_devices assert in config (#357)

* Remove n_devices assert * Update docstring
TransformerLensOrg · Jul 26, 2023 · 0316cc7 · 0316cc7
1 parent 03a5b52
commit 0316cc7
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 5 deletions.
diff --git a/transformer_lens/HookedTransformerConfig.py b/transformer_lens/HookedTransformerConfig.py
@@ -235,9 +235,6 @@ def __post_init__(self):
             self.device = utils.get_device()
 
         if self.n_devices > 1:
-            assert (
-                self.device == "cuda"
-            ), "n_devices > 1 is only supported on CUDA devices"
             assert (
                 torch.cuda.device_count() >= self.n_devices
             ), f"Not enough CUDA devices to support n_devices {self.n_devices}"

diff --git a/transformer_lens/utilities/devices.py b/transformer_lens/utilities/devices.py
@@ -47,8 +47,7 @@ def move_to_and_update_config(
     print_details=True,
 ):
     """
-    Wrapper around to that also changes model.cfg.device if it's a torch.device or string.
-    If torch.dtype, just passes through
+    Wrapper around `to` that also updates `model.cfg`.
     """
     if isinstance(device_or_dtype, torch.device):
         model.cfg.device = device_or_dtype.type