add mps to _set_device() in io

AdaptiveMotorControlLab · Jul 5, 2023 · 42965ea · 42965ea
1 parent e92f9bc
commit 42965ea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cebra/io.py b/cebra/io.py
@@ -75,7 +75,7 @@ def _set_device(self, device):
             return
         if not isinstance(device, str):
             device = device.type
-        if device not in ("cpu", "cuda"):
+        if device not in ("cpu", "cuda", "mps"):
             if device.startswith("cuda"):
                 _, id_ = device.split(":")
                 if int(id_) >= torch.cuda.device_count():