Suppress warning

sdv-dev · Jul 30, 2024 · bb3a087 · bb3a087
1 parent a0e0a76
commit bb3a087
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/sdv/single_table/ctgan.py b/sdv/single_table/ctgan.py
@@ -1,5 +1,7 @@
 """Wrapper around CTGAN model."""
 
+import warnings
+
 import numpy as np
 import pandas as pd
 import plotly.express as px
@@ -287,7 +289,9 @@ def _fit(self, processed_data):
             self.get_metadata(), processed_data, transformers
         )
         self._model = CTGAN(**self._model_kwargs)
-        self._model.fit(processed_data, discrete_columns=discrete_columns)
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', category=UserWarning)
+            self._model.fit(processed_data, discrete_columns=discrete_columns)
 
     def _sample(self, num_rows, conditions=None):
         """Sample the indicated number of rows from the model.