From f5967654fc1cc61c2c560a10871754f4142b8da9 Mon Sep 17 00:00:00 2001 From: Smiril Date: Tue, 12 Jul 2022 12:22:12 +0200 Subject: [PATCH] Update base_options.py --- options/base_options.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/options/base_options.py b/options/base_options.py index 4ee0209..35ba46e 100644 --- a/options/base_options.py +++ b/options/base_options.py @@ -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: + mps_device = torch.device("mps") args = vars(self.opt)