feat(wgan-gp): Add to init.

ydataai · Oct 29, 2020 · b5867da · b5867da
1 parent 5d37e62
commit b5867da
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/src/ydata_synthetic/synthesizers/__init__.py b/src/ydata_synthetic/synthesizers/__init__.py
@@ -1,7 +1,7 @@
 from ydata_synthetic.synthesizers.regular.cgan.model import CGAN
 from ydata_synthetic.synthesizers.regular.wgan.model import WGAN
 from ydata_synthetic.synthesizers.regular.vanillagan.model import VanilllaGAN
-
+from ydata_synthetic.synthesizers.regular.wgan_gp.model import WGAN_GP
 
 __all__ = [
     "VanilllaGAN",

diff --git a/src/ydata_synthetic/synthesizers/regular/wgan_gp/model.py b/src/ydata_synthetic/synthesizers/regular/wgan_gp/model.py
@@ -12,7 +12,6 @@
 from tensorflow.keras import Model
 from tensorflow.keras.optimizers import Adam
 
-
 #Auxiliary Keras backend class to calculate the Random Weighted average
 #https://stackoverflow.com/questions/58133430/how-to-substitute-keras-layers-merge-merge-in-tensorflow-keras
 class RandomWeightedAverage(tf.keras.layers.Layer):
@@ -27,7 +26,7 @@ def call(self, inputs, **kwargs):
     def compute_output_shape(self, input_shape):
         return input_shape[0]
 
-class WGAN(gan.Model):
+class WGAN_GP(gan.Model):
 
     def __init__(self, model_parameters, n_critic):
         # As recommended in WGAN paper - https://arxiv.org/abs/1701.07875