only load encoders of AEs

cvae.py

@@ -1,5 +1,6 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+from configs.config_cvae import config
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
 
 gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -35,24 +36,16 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.cvae_model import CVAE
-from configs.config_cvae import config
 
 
 def init_wandb():
     run = wandb.init(project='vae', name=config['name'],
-            notes='\
-offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
-retrained body encoder with deoriented bodies \
-use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\
-fixed gender issue and retrained body encoder, FINAL\
-remove shape from unposed garment\
-training on 100 garments with offsets over the average of these\
-fixed autoencoders!!!!!!!!!')
+            notes=config['notes'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('cvae.py')
     artifact.add_file('models/cvae_model.py')
     artifact.add_file('losses.py')
-    artifact.add_file('configs/config.py')
+    artifact.add_file('configs/config_base.py')
     artifact.add_file('configs/config_cvae.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
@@ -76,7 +69,7 @@ def kl_loss(mean, logvar):
     return kl_loss
 
 
-# @tf.function
+@tf.function
 def train_step(batch):
 
     with tf.GradientTape() as tape:
@@ -321,7 +314,7 @@ def sample(folder):
 
 
 if __name__ == '__main__':
-    # init_wandb()
+    init_wandb()
     allow_memory_growth()
     parser = argparse.ArgumentParser(description='')
     parser.add_argument('--name', type=str)