add GPU support on Mac

datamllab · Jun 29, 2023 · 05cc499 · 05cc499
1 parent a093f51
commit 05cc499
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/rlcard/utils/utils.py b/rlcard/utils/utils.py
@@ -19,7 +19,10 @@ def set_seed(seed):
 
 def get_device():
     import torch
-    if torch.cuda.is_available():
+    if torch.backends.mps.is_available():
+        device = torch.device("mps:0")
+        print("--> Running on the GPU")
+    elif torch.cuda.is_available():
         device = torch.device("cuda:0")
         print("--> Running on the GPU")
     else: