Repository: sdv-dev/CTGAN (language: python)

Issue:
Cannot save `CTGANSynthesizer` after sampling (`TypeError`)
### Environment Details

* SDV version: 1.0.0 (in progress)
* Python version: 3.8
* Operating System: Linux (Google Colab)

### Error Description
When using the `CTGANSynthesizer` from the new SDV 1.0 branch, I'm unable to save the trained synthesizer object.

### Steps to reproduce

```python
from sdv.datasets.demo import download_demo
from sdv.single_table import CTGANSynthesizer

real_data, metadata = download_demo(
    modality='single_table',
    dataset_name='fake_hotel_guests'
)

synthesizer = CTGANSynthesizer(metadata)
synthesizer.fit(real_data)
synthetic_data = synthesizer.sample(num_rows=500)
synthesizer.save('my_synthesizer.pkl')
```

### Stack Trace
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-10-70c571761350>](https://localhost:8080/#) in <module>
----> 1 synthesizer.save('my_synthesizer.pkl')
      2 
      3 # in the future
      4 synthesizer = CTGANSynthesizer.load('my_synthesizer.pkl')

2 frames
[/usr/local/lib/python3.8/dist-packages/sdv/single_table/base.py](https://localhost:8080/#) in save(self, filepath)
   1075         """
   1076         with open(filepath, 'wb') as output:
-> 1077             cloudpickle.dump(self, output)
   1078 
   1079     @classmethod

[/usr/local/lib/python3.8/dist-packages/cloudpickle/cloudpickle_fast.py](https://localhost:8080/#) in dump(obj, file, protocol, buffer_callback)
     53         compatibility with older versions of Python.
     54         """
---> 55         CloudPickler(
     56             file, protocol=protocol, buffer_callback=buffer_callback
     57         ).dump(obj)

[/usr/local/lib/python3.8/dist-packages/cloudpickle/cloudpickle_fast.py](https://localhost:8080/#) in dump(self, obj)
    630     def dump(self, obj):
    631         try:
--> 632             return Pickler.dump(self, obj)
    633         except RuntimeError as e:
    634             if "recursion" in e.args[0]:

TypeError: cannot pickle 'torch._C.Generator' object
```


Discussion / hints:

Fix patch:
diff --git a/ctgan/synthesizers/base.py b/ctgan/synthesizers/base.py
index e4ba7b82..0d31ef1c 100644
--- a/ctgan/synthesizers/base.py
+++ b/ctgan/synthesizers/base.py
@@ -57,13 +57,60 @@ def wrapper(self, *args, **kwargs):
 
 
 class BaseSynthesizer:
-    """Base class for all default synthesizers of ``CTGAN``.
-
-    This should contain the save/load methods.
-    """
+    """Base class for all default synthesizers of ``CTGAN``."""
 
     random_states = None
 
+    def __getstate__(self):
+        """Improve pickling state for ``BaseSynthesizer``.
+
+        Convert to ``cpu`` device before starting the pickling process in order to be able to
+        load the model even when used from an external tool such as ``SDV``. Also, if
+        ``random_states`` are set, store their states as dictionaries rather than generators.
+
+        Returns:
+            dict:
+                Python dict representing the object.
+        """
+        device_backup = self._device
+        self.set_device(torch.device('cpu'))
+        state = self.__dict__.copy()
+        self.set_device(device_backup)
+        if (
+            isinstance(self.random_states, tuple) and
+            isinstance(self.random_states[0], np.random.RandomState) and
+            isinstance(self.random_states[1], torch.Generator)
+        ):
+            state['_numpy_random_state'] = self.random_states[0].get_state()
+            state['_torch_random_state'] = self.random_states[1].get_state()
+            state.pop('random_states')
+
+        return state
+
+    def __setstate__(self, state):
+        """Restore the state of a ``BaseSynthesizer``.
+
+        Restore the ``random_states`` from the state dict if those are present and then
+        set the device according to the current hardware.
+        """
+        if '_numpy_random_state' in state and '_torch_random_state' in state:
+            np_state = state.pop('_numpy_random_state')
+            torch_state = state.pop('_torch_random_state')
+
+            current_torch_state = torch.Generator()
+            current_torch_state.set_state(torch_state)
+
+            current_numpy_state = np.random.RandomState()
+            current_numpy_state.set_state(np_state)
+            state['random_states'] = (
+                current_numpy_state,
+                current_torch_state
+            )
+
+        self.__dict__ = state
+        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.set_device(device)
+
     def save(self, path):
         """Save the model in the passed `path`."""
         device_backup = self._device


Test patch:
diff --git a/tests/integration/synthesizer/test_ctgan.py b/tests/integration/synthesizer/test_ctgan.py
index bdb8cc55..2faa1dc9 100644
--- a/tests/integration/synthesizer/test_ctgan.py
+++ b/tests/integration/synthesizer/test_ctgan.py
@@ -210,7 +210,7 @@ def test_fixed_random_seed():
     })
     discrete_columns = ['discrete']
 
-    ctgan = CTGAN(epochs=1)
+    ctgan = CTGAN(epochs=1, cuda=False)
 
     # Run
     ctgan.fit(data, discrete_columns)
@@ -273,3 +273,27 @@ def test_conditional():
 
 def test_batch_size_pack_size():
     """Test that if batch size is not a multiple of pack size, it raises a sane error."""
+
+
+def test_ctgan_save_and_load(tmpdir):
+    """Test that the ``CTGAN`` model can be saved and loaded."""
+    # Setup
+    data = pd.DataFrame({
+        'continuous': np.random.random(100),
+        'discrete': np.random.choice(['a', 'b', 'c'], 100)
+    })
+    discrete_columns = [1]
+
+    ctgan = CTGAN(epochs=1)
+    ctgan.fit(data.to_numpy(), discrete_columns)
+    ctgan.set_random_state(0)
+
+    ctgan.sample(100)
+    model_path = tmpdir / 'model.pkl'
+
+    # Save
+    ctgan.save(str(model_path))
+
+    # Load
+    loaded_instance = CTGAN.load(str(model_path))
+    loaded_instance.sample(100)
diff --git a/tests/integration/synthesizer/test_tvae.py b/tests/integration/synthesizer/test_tvae.py
index da7a6bc6..fdabe09f 100644
--- a/tests/integration/synthesizer/test_tvae.py
+++ b/tests/integration/synthesizer/test_tvae.py
@@ -129,3 +129,27 @@ def test_fixed_random_seed():
     assert not np.array_equal(sampled_random, sampled_0_1)
     np.testing.assert_array_equal(sampled_0_0, sampled_1_0)
     np.testing.assert_array_equal(sampled_0_1, sampled_1_1)
+
+
+def test_tvae_save(tmpdir):
+    """Test that the ``TVAE`` model can be saved and loaded."""
+    # Setup
+    data = pd.DataFrame({
+        'continuous': np.random.random(100),
+        'discrete': np.random.choice(['a', 'b', 'c'], 100)
+    })
+    discrete_columns = [1]
+
+    tvae = TVAE(epochs=1)
+    tvae.fit(data.to_numpy(), discrete_columns)
+    tvae.set_random_state(0)
+
+    tvae.sample(100)
+    model_path = tmpdir / 'model.pkl'
+
+    # Save
+    tvae.save(str(model_path))
+
+    # Load
+    loaded_instance = TVAE.load(str(model_path))
+    loaded_instance.sample(100)