working auto encoders

models/base.py

@@ -64,7 +64,10 @@ class AEEncoder(tf.keras.Model):
         x = self.encoder(x, training=training)
         x = self.bottleneck(x, training=training)
 
-        return x
+        if len(inputs) == 2:
+            return x, inputs[1]
+        else:
+            return x
 
 
 class Decoder(tf.keras.Model):

models/cond_encoder.py

@@ -8,7 +8,7 @@ class AutoEncoder(tf.keras.Model):
         super(AutoEncoder, self).__init__()
 
         self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
-        self.decoder : Decoder= Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
+        self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
 
 
     def call(self, inputs, training=None):