_get_enable_gpu_value + fix backward compatibility

Author: R-Palazzo

ctgan/synthesizers/_utils.py

@@ -4,7 +4,7 @@
 import torch
 
 
-def _validate_gpu_parameters(enable_gpu, cuda):
+def _get_enable_gpu_value(enable_gpu, cuda):
     """Validate both the `enable_gpu` and `cuda` parameters.
 
     The logic here is to:
@@ -53,5 +53,5 @@ def _set_device(enable_gpu, device=None):
 
 
 def validate_and_set_device(enable_gpu, cuda):
-    enable_gpu = _validate_gpu_parameters(enable_gpu, cuda)
+    enable_gpu = _get_enable_gpu_value(enable_gpu, cuda)
     return _set_device(enable_gpu)

ctgan/synthesizers/ctgan.py

@@ -541,6 +541,7 @@ def sample(self, n, condition_column=None, condition_value=None):
 
     def set_device(self, device):
         """Set the `device` to be used ('GPU' or 'CPU)."""
-        self._device = _set_device(self._enable_gpu, device)
+        enable_gpu = getattr(self, '_enable_gpu', True)
+        self._device = _set_device(enable_gpu, device)
         if self._generator is not None:
             self._generator.to(self._device)

ctgan/synthesizers/tvae.py

@@ -237,5 +237,6 @@ def sample(self, samples):
 
     def set_device(self, device):
         """Set the `device` to be used ('GPU' or 'CPU')."""
-        self._device = _set_device(self._enable_gpu, device)
+        enable_gpu = getattr(self, '_enable_gpu', True)
+        self._device = _set_device(enable_gpu, device)
         self.decoder.to(self._device)

tests/unit/synthesizer/test__utils.py

@@ -5,11 +5,11 @@
 import pytest
 import torch
 
-from ctgan.synthesizers._utils import _set_device, _validate_gpu_parameters, validate_and_set_device
+from ctgan.synthesizers._utils import _get_enable_gpu_value, _set_device, validate_and_set_device
 
 
 def test__validate_gpu_parameter():
-    """Test the ``_validate_gpu_parameters`` method."""
+    """Test the ``_get_enable_gpu_value`` method."""
     # Setup
     expected_error = re.escape(
         'Cannot resolve the provided values of `cuda` and `enable_gpu` parameters. '
@@ -21,13 +21,13 @@ def test__validate_gpu_parameter():
     )
 
     # Run
-    enable_gpu_1 = _validate_gpu_parameters(False, None)
-    enable_gpu_2 = _validate_gpu_parameters(True, None)
+    enable_gpu_1 = _get_enable_gpu_value(False, None)
+    enable_gpu_2 = _get_enable_gpu_value(True, None)
     with pytest.warns(FutureWarning, match=expected_warning):
-        enable_gpu_3 = _validate_gpu_parameters(True, False)
+        enable_gpu_3 = _get_enable_gpu_value(True, False)
 
     with pytest.raises(ValueError, match=expected_error):
-        _validate_gpu_parameters(False, True)
+        _get_enable_gpu_value(False, True)
 
     # Assert
     assert enable_gpu_1 is False
@@ -62,7 +62,7 @@ def test__set_device():
 
 
 @patch('ctgan.synthesizers._utils._set_device')
-@patch('ctgan.synthesizers._utils._validate_gpu_parameters')
+@patch('ctgan.synthesizers._utils._get_enable_gpu_value')
 def test_validate_and_set_device(mock_validate, mock_set_device):
     """Test the ``validate_and_set_device`` method."""
     # Setup