Merge pull request #83 from ahoust17/master

pycroscopy · Jan 8, 2024 · 06b121c · 06b121c
2 parents e884ba5 + 0802e4a
commit 06b121c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/atomai/trainers/trainer.py b/atomai/trainers/trainer.py
@@ -67,8 +67,8 @@ class BaseTrainer:
     """
     def __init__(self):
         set_train_rng(1)
-        if torch.backends.mps.is_available():
-            self.device = torch.device("mps") # backend for Apple silicon GPUs
+        if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            self.device = torch.device('mps') # backend for Apple silicon GPUs
         elif torch.cuda.is_available():
             self.device = 'cuda'
         else:
@@ -367,8 +367,8 @@ def print_statistics(self, e: int, **kwargs) -> None:
             accuracy_metrics = "Accuracy"
         if torch.cuda.is_available():
             gpu_usage = gpu_usage_map(torch.cuda.current_device())
-        elif torch.backends.mps.is_available():
-            gpu_usage = gpu_usage_map(torch.mps.current_device())
+        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            gpu_usage = ['N/A ', ' N/A']
         else:
             gpu_usage = ['N/A ', ' N/A']
         if self.compute_accuracy:

diff --git a/atomai/utils/nn.py b/atomai/utils/nn.py
@@ -48,7 +48,7 @@ def load_weights(model: Type[torch.nn.Module],
     torch.manual_seed(0)
     if torch.cuda.device_count() > 0:
         checkpoint = torch.load(weights_path)
-    elif torch.backends.mps.is_available():
+    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
         checkpoint = torch.load(weights_path, map_location='mps')
     else:
         checkpoint = torch.load(weights_path, map_location='cpu')