Skip to content

Commit

Permalink
Check if decoder is set
Browse files Browse the repository at this point in the history
Without this check things like the sklearn.randomsampler run into an error, since they deepcopy the class before they call the fit method
  • Loading branch information
amarcelq authored Oct 25, 2024
1 parent bc8235c commit a04add4
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion ctgan/synthesizers/tvae.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,8 @@ def __init__(
self.loss_values = pd.DataFrame(columns=['Epoch', 'Batch', 'Loss'])
self.verbose = verbose

self.decoder = None

if not cuda or not torch.cuda.is_available():
device = 'cpu'
elif isinstance(cuda, str):
Expand Down Expand Up @@ -243,4 +245,5 @@ def sample(self, samples):
def set_device(self, device):
"""Set the `device` to be used ('GPU' or 'CPU)."""
self._device = device
self.decoder.to(self._device)
if self.decoder:
self.decoder.to(self._device)

0 comments on commit a04add4

Please sign in to comment.