Fork the CUDA RNG explicitly in deterministic_rng

tests/unit/benchmarks/test_deterministic_rng.py

@@ -21,6 +21,10 @@ def setter(name):
     monkeypatch.setattr(torch.cuda, "manual_seed", setter("cuda"))
     monkeypatch.setattr(torch.cuda, "manual_seed_all", setter("cuda"))
     monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
+    # As on Apple Silicon, MPS is the current accelerator
+    monkeypatch.setattr(
+        torch.accelerator, "current_accelerator", lambda *args, **kwargs: torch.device("mps")
+    )
     monkeypatch.setattr(torch.mps, "get_rng_state", lambda *args, **kwargs: state["mps"])
     monkeypatch.setattr(torch.mps, "set_rng_state", setter("mps"))
     monkeypatch.setattr(torch.mps, "manual_seed", setter("mps"))

transformer_lens/benchmarks/utils.py

@@ -27,7 +27,8 @@ def deterministic_rng(seed: int = 42) -> Iterator[None]:
     # generators snapshotted here are seeded, so each one is restored afterwards.
     cuda_devices = list(range(torch.cuda.device_count())) if torch.cuda.is_available() else []
     mps_state = torch.mps.get_rng_state() if torch.backends.mps.is_available() else None
-    with torch.random.fork_rng(devices=cuda_devices):
+    # fork_rng defaults to the current accelerator since torch 2.13, so name CUDA explicitly.
+    with torch.random.fork_rng(devices=cuda_devices, device_type="cuda"):
         try:
             torch.default_generator.manual_seed(seed)
             if cuda_devices: