add testing to static ae

models/base.py

@@ -1,7 +1,4 @@
-import os
 import tensorflow as tf
-
-from custom_utils.custom_utils import load_model
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
 
 class Bottleneck(tf.keras.Model):
@@ -13,7 +10,6 @@ class Bottleneck(tf.keras.Model):
                 [Block(dim=dim, drop_path=0, kernel_size=kernel_size,
                        layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
             )
-        # self.bottleneck = [Block(dim=dim, drop_path=0, kernel_size=kernel_size, layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
             
         self.flatten = tf.keras.layers.Flatten()
         self.dense = tf.keras.layers.Dense(latent_dim)
@@ -21,9 +17,7 @@ class Bottleneck(tf.keras.Model):
     def call(self, inputs, training=None):
         x = inputs
 
-        # for bt in self.bottleneck:
-        #     x = bt(x, training=training)
-        x = self.bottleneck(x)
+        x = self.bottleneck(x, training=training)
         x = self.flatten(x)
         x = self.dense(x)
 
@@ -32,7 +26,7 @@ class Bottleneck(tf.keras.Model):
 
 class VariationalEncoder(tf.keras.Model):
     
-    def __init__(self, latent_dim, depths, dims, use_bottleneck=True, kernel_size=7):
+    def __init__(self, latent_dim, depths, dims, kernel_size=7):
         super(VariationalEncoder, self).__init__()
 
         self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)