TVAE: Adapt batch_size to data size

[csala]

Environment Details

Please indicate the following details about the environment in which you found the bug:

• CTGAN version: 0.4.0

Error Description

Current TVAE fails to be fitted on datasets that are shorter than the batch size (and it does it silently).

The culprit is the drop_last=True in this line:

CTGAN/ctgan/synthesizers/tvae.py

Line 113 in 38f0d30

loader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, drop_last=True)

As a consequence of this line, if the data is shorter than the batch_size TVAE does nothing during the training and results in a model that kept the initialization weights, unoptimized.

Fix

The indicated line could be changed to something like this which:

• Reduces the batch size to the data length if the data has less rows than the batch size
• Drops the last batch only if the data size is not divisible by the batch size

        data_length = len(table_data)
        batch_size = min(self.batch_size, data_length)
        drop_last = (data_length % batch_size) != 0
        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=drop_last)