define _set_device()

Author: R-Palazzo

deepecho/models/basic_gan.py

@@ -14,6 +14,34 @@
 LOGGER = logging.getLogger(__name__)
 
 
+def _set_device(enable_gpu, cuda):
+    if cuda is not None:
+        if not enable_gpu:
+            raise ValueError(
+                'Cannot set `cuda` and `enable_gpu` together. Please use only `enable_gpu`.'
+            )
+
+        warnings.warn(
+            '`cuda` parameter is deprecated and will be removed in a future release. '
+            'Please use `enable_gpu` instead.',
+            FutureWarning,
+        )
+        enable_gpu = cuda
+
+    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'
+
+    return torch.device(device)
+
+
 def _expand_context(data, context):
     return torch.cat(
         [
@@ -165,36 +193,16 @@ def __init__(
         gen_lr=1e-3,
         dis_lr=1e-3,
         enable_gpu=True,
-        cuda=None,
         verbose=True,
+        cuda=None,
     ):
-        if cuda is not None:
-            warnings.warn(
-                '`cuda` parameter is deprecated and will be removed in a future release. '
-                'Please use `enable_gpu` instead.',
-                FutureWarning,
-            )
-            enable_gpu = cuda
-
         self._epochs = epochs
         self._gen_lr = gen_lr
         self._dis_lr = dis_lr
         self._latent_size = latent_size
         self._hidden_size = hidden_size
-        self._enable_gpu = enable_gpu
-
-        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'
-
-        self._device = torch.device(device)
+        self._device = _set_device(enable_gpu, cuda)
+        self._enable_gpu = cuda if cuda is not None else enable_gpu
         self._verbose = verbose
 
         LOGGER.info('%s instance created', self)

tests/integration/test_basic_gan.py

@@ -7,7 +7,44 @@
 import pytest
 import torch
 
-from deepecho.models.basic_gan import BasicGANModel
+from deepecho.models.basic_gan import BasicGANModel, _set_device
+
+
+def test__set_device():
+    """Test the ``_set_device`` method."""
+    # Setup
+    expected_error = re.escape(
+        'Cannot set `cuda` and `enable_gpu` together. Please use only `enable_gpu`.'
+    )
+    expected_warning = re.escape(
+        '`cuda` parameter is deprecated and will be removed in a future release. '
+        'Please use `enable_gpu` instead.'
+    )
+
+    # Run
+    device_1 = _set_device(False, None)
+    device_2 = _set_device(True, None)
+    with pytest.warns(FutureWarning, match=expected_warning):
+        device_3 = _set_device(True, False)
+
+    with pytest.raises(ValueError, match=expected_error):
+        _set_device(False, True)
+
+    # Assert
+    if (
+        sys.platform == 'darwin'
+        and getattr(torch.backends, 'mps', None)
+        and torch.backends.mps.is_available()
+    ):
+        expected_device_2 = torch.device('mps')
+    elif torch.cuda.is_available():
+        expected_device_2 = torch.device('cuda')
+    else:
+        expected_device_2 = torch.device('cpu')
+
+    assert device_1 == torch.device('cpu')
+    assert device_2 == expected_device_2
+    assert device_3 == torch.device('cpu')
 
 
 class TestBasicGANModel(unittest.TestCase):