test without env variable

Author: R-Palazzo

deepecho/__init__.py

@@ -4,6 +4,11 @@
 __email__ = '[email protected]'
 __version__ = '0.7.1.dev0'
 __path__ = __import__('pkgutil').extend_path(__path__, __name__)
+import os
+
+# Enable fallback so ops not implemented on MPS run on CPU
+# https://github.com/pytorch/pytorch/issues/77764
+os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
 
 from deepecho.demo import load_demo
 from deepecho.models.basic_gan import BasicGANModel

deepecho/models/_utils.py

@@ -1,4 +1,10 @@
-import sys
+import os
+
+# Enable fallback so ops not implemented on MPS run on CPU
+# https://github.com/pytorch/pytorch/issues/77764
+os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
+
+import platform
 import warnings
 
 import torch
@@ -32,15 +38,16 @@ def _validate_gpu_parameters(enable_gpu, cuda):
 def _set_device(enable_gpu):
     """Set the torch device based on the `enable_gpu` parameter and system capabilities."""
     if enable_gpu:
-        if sys.platform == 'darwin':  # macOS
-            if getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
-                device = 'mps'
-            else:
-                device = 'cpu'
-        else:  # Linux/Windows
-            device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    else:
-        device = 'cpu'
+        if (
+            platform.machine() == 'arm64'
+            and getattr(torch.backends, 'mps', None)
+            and torch.backends.mps.is_available()
+        ):
+            device = 'mps'
+        else:
+            device = 'cpu'
+    else:  # Linux/Windows
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     return torch.device(device)
 

deepecho/models/par.py

@@ -7,6 +7,7 @@
 import torch
 from tqdm import tqdm
 
+from deepecho.models._utils import validate_and_set_device
 from deepecho.models.base import DeepEcho
 
 LOGGER = logging.getLogger(__name__)
@@ -98,18 +99,11 @@ class PARModel(DeepEcho):
             Whether to print progress to console or not.
     """
 
-    def __init__(self, epochs=128, sample_size=1, cuda=True, verbose=True):
+    def __init__(self, epochs=128, sample_size=1, enable_gpu=True, verbose=True, cuda=None):
         self.epochs = epochs
         self.sample_size = sample_size
-
-        if not cuda or not torch.cuda.is_available():
-            device = 'cpu'
-        elif isinstance(cuda, str):
-            device = cuda
-        else:
-            device = 'cuda'
-
-        self.device = torch.device(device)
+        self.device = validate_and_set_device(enable_gpu=enable_gpu, cuda=cuda)
+        self._enable_gpu = cuda if cuda is not None else enable_gpu
         self.verbose = verbose
         self.loss_values = pd.DataFrame(columns=['Epoch', 'Loss'])