Add M1 GPU Support

not only CUDA Support by adding lines
2022-09-24 20:35:33 +02:00
parent dd1ecdd2a7
commit 2f2d68e04a
1 changed files with 10 additions and 2 deletions
@@ -80,8 +80,16 @@ class BaseOptions():
                self.opt.gpu_ids.append(id)
        
        # set gpu ids
-        if len(self.opt.gpu_ids) > 0:
-            torch.cuda.set_device(self.opt.gpu_ids[0])
+        if not torch.backends.mps.is_available():
+           if not torch.backends.mps.is_built():
+             if len(self.opt.gpu_ids) > 0:
+               torch.cuda.set_device(self.opt.gpu_ids[0])
+             else:
+                print("ERROR")
+           else:
+             print("ERROR")
+        else:
+          torch.device("mps")

        args = vars(self.opt)