Skip to content

Commit

Permalink
Suppress warning when using GANs that suggests that CUDA isn't being …
Browse files Browse the repository at this point in the history
…used (#2155)
  • Loading branch information
fealho authored Aug 7, 2024
1 parent 17aaa28 commit 1c07c22
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion sdv/single_table/ctgan.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
"""Wrapper around CTGAN model."""

import warnings

import numpy as np
import pandas as pd
import plotly.express as px
Expand Down Expand Up @@ -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.filterwarnings('ignore', message='.*Attempting to run cuBLAS.*')
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.
Expand Down

0 comments on commit 1c07c22

Please sign in to comment.