Verify err msg matches

sdv-dev · May 20, 2024 · f483afb · f483afb
1 parent 7749f05
commit f483afb
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 12 deletions.
diff --git a/sdv/multi_table/base.py b/sdv/multi_table/base.py
@@ -696,12 +696,19 @@ def load(cls, filepath):
         with open(filepath, 'rb') as f:
             try:
                 synthesizer = cloudpickle.load(f)
-            except RuntimeError:
-                raise SamplingError(
-                    'This synthesizer was created on a machine with GPU but the current machine is'
-                    ' CPU-only. This feature is currently unsupported. We recommend sampling on '
-                    'the same GPU-enabled machine.'
+            except RuntimeError as e:
+                err_msg = (
+                    'Attempting to deserialize object on a CUDA device but '
+                    'torch.cuda.is_available() is False. If you are running on a CPU-only machine,'
+                    " please use torch.load with map_location=torch.device('cpu') "
+                    'to map your storages to the CPU.'
                 )
+                if str(e) == err_msg:
+                    raise SamplingError(
+                        'This synthesizer was created on a machine with GPU but the current machine is'
+                        ' CPU-only. This feature is currently unsupported. We recommend sampling on '
+                        'the same GPU-enabled machine.'
+                    )
 
         check_synthesizer_version(synthesizer)
         check_sdv_versions_and_warn(synthesizer)

diff --git a/sdv/single_table/base.py b/sdv/single_table/base.py
@@ -494,12 +494,19 @@ def load(cls, filepath):
         with open(filepath, 'rb') as f:
             try:
                 synthesizer = cloudpickle.load(f)
-            except RuntimeError:
-                raise SamplingError(
-                    'This synthesizer was created on a machine with GPU but the current machine is'
-                    ' CPU-only. This feature is currently unsupported. We recommend sampling on '
-                    'the same GPU-enabled machine.'
+            except RuntimeError as e:
+                err_msg = (
+                    'Attempting to deserialize object on a CUDA device but '
+                    'torch.cuda.is_available() is False. If you are running on a CPU-only machine,'
+                    " please use torch.load with map_location=torch.device('cpu') "
+                    'to map your storages to the CPU.'
                 )
+                if str(e) == err_msg:
+                    raise SamplingError(
+                        'This synthesizer was created on a machine with GPU but the current machine is'
+                        ' CPU-only. This feature is currently unsupported. We recommend sampling on '
+                        'the same GPU-enabled machine.'
+                    )
 
         check_synthesizer_version(synthesizer)
         check_sdv_versions_and_warn(synthesizer)

diff --git a/tests/unit/multi_table/test_base.py b/tests/unit/multi_table/test_base.py
@@ -1609,7 +1609,12 @@ def test_load(self, mock_file, cloudpickle_mock,
     def test_load_runtime_error(self, cloudpickle_mock, mock_open):
         """Test that the synthesizer's load method errors with the correct message."""
         # Setup
-        cloudpickle_mock.load.side_effect = RuntimeError
+        cloudpickle_mock.load.side_effect = RuntimeError((
+            'Attempting to deserialize object on a CUDA device but '
+            'torch.cuda.is_available() is False. If you are running on a CPU-only machine,'
+            " please use torch.load with map_location=torch.device('cpu') "
+            'to map your storages to the CPU.'
+        ))
 
         # Run and Assert
         err_msg = re.escape(

diff --git a/tests/unit/single_table/test_base.py b/tests/unit/single_table/test_base.py
@@ -1919,7 +1919,12 @@ def test_load_custom_constraint_classes(self):
     def test_load_runtime_error(self, cloudpickle_mock, mock_open):
         """Test that the synthesizer's load method errors with the correct message."""
         # Setup
-        cloudpickle_mock.load.side_effect = RuntimeError
+        cloudpickle_mock.load.side_effect = RuntimeError((
+            'Attempting to deserialize object on a CUDA device but '
+            'torch.cuda.is_available() is False. If you are running on a CPU-only machine,'
+            " please use torch.load with map_location=torch.device('cpu') "
+            'to map your storages to the CPU.'
+        ))
 
         # Run and Assert
         err_msg = re.escape(