Compatibility with Apple Silicon GPUs

atomai/__version__.py

@@ -1 +1 @@
-version = '0.7.8'
+version = '0.7.8.MAC'

atomai/trainers/trainer.py

@@ -67,7 +67,12 @@ class BaseTrainer:
     """
     def __init__(self):
         set_train_rng(1)
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        if torch.backends.mps.is_available():
+            self.device = torch.device("mps") # backend for Apple silicon GPUs
+        elif torch.cuda.is_available():
+            self.device = 'cuda'
+        else:
+            self.device = 'cpu'
         self.net = None
         self.criterion = None
         self.optimizer = None
@@ -362,6 +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())
         else:
             gpu_usage = ['N/A ', ' N/A']
         if self.compute_accuracy:

atomai/utils/nn.py

@@ -48,6 +48,8 @@ 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():
+        checkpoint = torch.load(weights_path, map_location='mps')
     else:
         checkpoint = torch.load(weights_path, map_location='cpu')
     model.load_state_dict(checkpoint)
@@ -199,6 +201,9 @@ def mock_forward(model: Type[torch.nn.Module],
     x = torch.randn(1, *dims)
     if next(model.parameters()).is_cuda:
         x = x.cuda()
+    elif next(model.parameters()).is_mps:
+        mps_device = torch.device("mps")
+        x = x.to(mps_device)
     out = model(x)
     return out